#pragma once
#include <cmath>
#include <algorithm>

enum class DistortionType { SoftClip = 0, HardClip, Foldback, Overdrive, NumTypes };

class Distortion {
public:
    void prepare(double /*sr*/) {}

    void setParameters(DistortionType type, float amount, float mix) {
        distType = type;
        distAmount = amount;
        dryWet = mix;
    }

    float process(float input) {
        float driven = input * (1.0f + distAmount * 9.0f);
        float wet = 0.0f;

        switch (distType) {
            case DistortionType::SoftClip:
                wet = std::tanh(driven);
                break;

            case DistortionType::HardClip: {
                float threshold = 1.0f / (1.0f + distAmount * 9.0f);
                wet = std::clamp(driven, -threshold, threshold);
                break;
            }

            case DistortionType::Foldback: {
                float threshold = 1.0f + distAmount * 3.0f;
                wet = driven;
                while (wet > threshold || wet < -threshold) {
                    if (wet > threshold)
                        wet = 2.0f * threshold - wet;
                    else if (wet < -threshold)
                        wet = -2.0f * threshold - wet;
                    else
                        break;
                }
                if (threshold > 0.0f)
                    wet /= threshold;
                break;
            }

            case DistortionType::Overdrive: {
                float t = driven;
                if (t > 1.0f) t = 1.0f;
                else if (t < -1.0f) t = -1.0f;
                wet = t * (1.5f - 0.5f * t * t);
                break;
            }

            default:
                wet = driven;
                break;
        }

        if (distAmount < 0.001f)
            return input;

        return input + (wet - input) * dryWet;
    }

private:
    DistortionType distType = DistortionType::SoftClip;
    float distAmount = 0.0f;
    float dryWet = 0.0f;
};