import Foundation

/// Length prior over candidate segment lengths. Port of `create_prior_function`
/// in `wtpsplit/utils/priors.py`. A prior returns a non-negative weight for a
/// given segment length; `0` forbids that length.
public enum Prior: String, Sendable, CaseIterable {
    case uniform
    case gaussian
    case clippedPolynomial
}

enum PriorFactory {
    /// Build the base prior closure. `targetLength`/`spread` are ignored by
    /// `.uniform` (matching the reference, which omits them for uniform).
    static func make(_ kind: Prior, maxLength: Int, targetLength: Int, spread: Int) -> (Int) -> Double {
        let maxL = maxLength
        switch kind {
        case .uniform:
            return { length in length > maxL ? 0.0 : 1.0 }

        case .gaussian:
            let target = Double(targetLength)
            let sp = Double(spread)
            return { length in
                if length > maxL { return 0.0 }
                let z = (Double(length) - target) / sp
                return exp(-0.5 * z * z)
            }

        case .clippedPolynomial:
            let target = Double(targetLength)
            let falloff = 1.0 / (Double(spread) * Double(spread))
            return { length in
                if length > maxL { return 0.0 }
                let d = Double(length) - target
                return max(1.0 - falloff * d * d, 0.0)
            }
        }
    }

    /// Wrap a base prior with the soft-overflow decay used by
    /// `run_segmentation.py`: lengths past `softCap` are damped by a Gaussian
    /// tail so plain word gaps can't exploit the slack, while a strong pause can
    /// still pull a break into the overflow zone. `overflow == 0` ⇒ hard cap.
    static func withOverflow(_ base: @escaping (Int) -> Double, softCap: Int, overflow: Int) -> (Int) -> Double {
        guard overflow > 0 else { return base }
        let decay = Double(overflow)
        return { length in
            let b = base(length)
            if length <= softCap { return b }
            let t = (Double(length - softCap)) / decay
            return b * exp(-(t * t))
        }
    }
}