<?php

declare(strict_types=1);

namespace Dochia;

/**
 * Two-component lognormal mixture.
 *
 * With moment_match = false (the shipped configuration) the body lognormal
 * keeps the target mean/sd and a right-tail component is added; the mixture
 * mean/sd therefore rise slightly, which is intentional: it protects the 80%
 * body interval and adds honest tail mass for long/skipped cycles.
 */
final class LognormalMixture
{
    /**
     * @param  list<array{0: float, 1: float, 2: float}>  $components  [weight, log-mean, log-sd]
     */
    public function __construct(
        public readonly float $mean,
        public readonly float $sd,
        public readonly float $tailProb,
        public readonly array $components,
    ) {
    }

    public static function fromMoments(
        float $mean,
        float $sd,
        float $pi,
        float $tailShift,
        float $tailSdMultiplier,
        bool $momentMatch = false,
    ): self {
        $pi = min(max($pi, 0.0), 0.45);
        $mean = max($mean, 1.0);
        $sd = max($sd, 0.2);

        if ($pi <= 1e-9) {
            [$m, $s] = self::lognormalParams($mean, $sd);
            return new self($mean, $sd, 0.0, [[1.0, $m, $s]]);
        }

        $tailMean = max($mean + $tailShift, $mean * 1.02);
        $tailSd = max($sd * $tailSdMultiplier, 0.3);

        if (! $momentMatch) {
            [$mMain, $sMain] = self::lognormalParams($mean, $sd);
            [$mTail, $sTail] = self::lognormalParams($tailMean, $tailSd);
            return new self($mean, $sd, $pi, [[1.0 - $pi, $mMain, $sMain], [$pi, $mTail, $sTail]]);
        }

        $mainMean = ($mean - $pi * $tailMean) / (1.0 - $pi);
        $mainVar = (($sd * $sd + $mean * $mean - $pi * ($tailSd * $tailSd + $tailMean * $tailMean)) / (1.0 - $pi))
            - $mainMean * $mainMean;
        if ($mainMean <= 0.5 || $mainVar <= 0.04) {
            [$m, $s] = self::lognormalParams($mean, $sd);
            return new self($mean, $sd, 0.0, [[1.0, $m, $s]]);
        }
        [$mMain, $sMain] = self::lognormalParams($mainMean, sqrt($mainVar));
        [$mTail, $sTail] = self::lognormalParams($tailMean, $tailSd);
        return new self($mean, $sd, $pi, [[1.0 - $pi, $mMain, $sMain], [$pi, $mTail, $sTail]]);
    }

    /** @return array{0: float, 1: float} */
    private static function lognormalParams(float $mean, float $sd): array
    {
        $mean = max($mean, 1e-6);
        $sd = max($sd, 1e-6);
        $s2 = log(1.0 + ($sd * $sd) / ($mean * $mean));
        $s = sqrt($s2);
        $m = log($mean) - 0.5 * $s2;
        return [$m, $s];
    }

    public function cdf(float $x): float
    {
        if ($x <= 0.0) {
            return 0.0;
        }
        $total = 0.0;
        foreach ($this->components as [$w, $m, $s]) {
            $total += $w * Erf::normCdf((log($x) - $m) / $s);
        }
        return min(max($total, 0.0), 1.0);
    }

    public function pdf(float $x): float
    {
        if ($x <= 0.0) {
            return 0.0;
        }
        $total = 0.0;
        foreach ($this->components as [$w, $m, $s]) {
            $z = (log($x) - $m) / $s;
            $total += $w * Erf::normPdf($z) / ($x * $s);
        }
        return $total;
    }

    public function quantile(float $q): float
    {
        if ($q <= 0.0) {
            return 0.0;
        }
        if ($q >= 1.0) {
            return 1000.0;
        }
        $lo = 1.0;
        $hi = 400.0;
        if ($this->cdf($lo) >= $q) {
            return $lo;
        }
        if ($this->cdf($hi) <= $q) {
            return $hi;
        }
        for ($i = 0; $i < 40; $i++) {
            $mid = 0.5 * ($lo + $hi);
            if ($this->cdf($mid) < $q) {
                $lo = $mid;
            } else {
                $hi = $mid;
            }
        }
        return 0.5 * ($lo + $hi);
    }
}