using System.Drawing;
using Emgu.CV;
using Emgu.CV.CvEnum;
using Emgu.CV.Structure;
using Emgu.CV.Util;
using XP.Calibration.Core;
using XP.Calibration.Models;

namespace XP.Calibration;

/// <summary>
/// 完整几何校准: TIGRE 投影模型 + LM 优化
/// Full geometric calibration: TIGRE projection model + Levenberg-Marquardt optimization
/// </summary>
public class FullCalibration
{
    // 参数索引常量
    private const int IDX_AY = 0;
    private const int IDX_DET_Y = 1;
    private const int IDX_DET_X = 2;
    private const int IDX_DET_Z = 3;
    private const int IDX_R = 4;
    private const int BASE_PARAMS = 5;

    // 参数边界
    private static readonly double[] LowerBase = { -Math.PI / 2, -Math.PI / 3, -Math.PI / 3, -Math.PI, 0.1 };
    private static readonly double[] UpperBase = { Math.PI / 2, Math.PI / 3, Math.PI / 3, Math.PI, 100 };

    /// <summary>
    /// 进度回调 | Progress callback (iteration, rms, params)
    /// </summary>
    public Action<int, double, double[]>? OnProgress { get; set; }

    /// <summary>
    /// 从投影序列执行完整几何校准
    /// </summary>
    public FullCalibrationResult Calibrate(
        List<Mat> projections, GeoParams geo, FullCalibrationOptions? options = null)
    {
        options ??= new FullCalibrationOptions();

        // 检测各帧球心
        var thetas = new List<double>();
        var uObs = new List<double>();
        var vObs = new List<double>();
        var pts = new List<PointF>();

        for (int i = 0; i < projections.Count; i++)
        {
            if (BallDetector.DetectCenter(projections[i], out var c))
            {
                thetas.Add((double)i / (projections.Count - 1) * 2.0 * Math.PI);
                uObs.Add(c.X);
                vObs.Add(c.Y);
                pts.Add(c);
            }
        }

        if (pts.Count <= 10)
            throw new InvalidOperationException($"有效检测帧数不足: {pts.Count}");

        // 椭圆拟合估算初值
        using var vp = new VectorOfPointF(pts.ToArray());
        var ell = CvInvoke.FitEllipse(vp);
        double longAxis = Math.Max(ell.Size.Width, ell.Size.Height);
        double shortAxis = Math.Min(ell.Size.Width, ell.Size.Height);
        double mag = geo.DSD / geo.DSO;
        double rInit = longAxis * geo.PixelSize / (2.0 * mag);
        double ayInit = Math.Acos(Math.Min(1.0, shortAxis / longAxis));

        // 确定参数个数和索引
        int np = BASE_PARAMS;
        int idxOffDU = -1, idxOffDV = -1, idxDSO = -1, idxDSD = -1;

        if (options.OptimizeDetectorOffset)
        {
            idxOffDU = np; idxOffDV = np + 1; np += 2;
        }
        if (options.OptimizeDistances)
        {
            idxDSO = np; idxDSD = np + 1; np += 2;
        }

        // 构造初始参数
        var p = new double[np];
        p[IDX_AY] = ayInit;
        p[IDX_DET_Y] = 0;
        p[IDX_DET_X] = 0;
        p[IDX_DET_Z] = 0;
        p[IDX_R] = rInit;
        if (options.OptimizeDetectorOffset)
        {
            p[idxOffDU] = 0; p[idxOffDV] = 0;
        }
        if (options.OptimizeDistances)
        {
            p[idxDSO] = geo.DSO; p[idxDSD] = geo.DSD;
        }

        // LM 优化
        var ctx = new LmContext(thetas, uObs, vObs, geo, np,
            options.OptimizeDetectorOffset, options.OptimizeDistances,
            idxOffDU, idxOffDV, idxDSO, idxDSD);

        bool converged = SolveLM(p, ctx, options.MaxIterations, options.Tolerance);

        // 计算最终 RMS
        var finalRes = ComputeResiduals(p, ctx);
        double rms = Math.Sqrt(finalRes.Sum(r => r * r) / finalRes.Length);

        double dsoOut = options.OptimizeDistances ? p[idxDSO] : geo.DSO;
        double dsdOut = options.OptimizeDistances ? p[idxDSD] : geo.DSD;

        return new FullCalibrationResult
        {
            AyDeg = p[IDX_AY] * 180.0 / Math.PI,
            DetYDeg = p[IDX_DET_Y] * 180.0 / Math.PI,
            DetXDeg = p[IDX_DET_X] * 180.0 / Math.PI,
            DetZDeg = p[IDX_DET_Z] * 180.0 / Math.PI,
            R_mm = p[IDX_R],
            OffDetecU_mm = options.OptimizeDetectorOffset ? p[idxOffDU] : 0,
            OffDetecV_mm = options.OptimizeDetectorOffset ? p[idxOffDV] : 0,
            DSO_mm = dsoOut,
            DSD_mm = dsdOut,
            RmsPx = rms,
            Converged = converged
        };
    }

    /// <summary>
    /// 从 RAW 文件执行完整几何校准 (便捷方法)
    /// </summary>
    public FullCalibrationResult CalibrateFromRaw(
        string rawPath, int width, int height, int count,
        GeoParams geo, FullCalibrationOptions? options = null)
    {
        var projections = RawReader.ReadFloat32(rawPath, width, height, count);
        try
        {
            return Calibrate(projections, geo, options);
        }
        finally
        {
            foreach (var p in projections) p.Dispose();
        }
    }

    #region LM Solver

    private record LmContext(
        List<double> Thetas, List<double> UObs, List<double> VObs,
        GeoParams Geo, int NP,
        bool OptOffset, bool OptDist,
        int IdxOffDU, int IdxOffDV, int IdxDSO, int IdxDSD);

    private static double[] ComputeResiduals(double[] p, LmContext ctx)
    {
        double ay = p[IDX_AY], detY = p[IDX_DET_Y];
        double detX = p[IDX_DET_X], detZ = p[IDX_DET_Z], R = p[IDX_R];
        double offU = ctx.OptOffset ? p[ctx.IdxOffDU] : 0;
        double offV = ctx.OptOffset ? p[ctx.IdxOffDV] : 0;

        var gp = new GeoParams
        {
            DSO = ctx.OptDist ? p[ctx.IdxDSO] : ctx.Geo.DSO,
            DSD = ctx.OptDist ? p[ctx.IdxDSD] : ctx.Geo.DSD,
            PixelSize = ctx.Geo.PixelSize,
            NDetecU = ctx.Geo.NDetecU,
            NDetecV = ctx.Geo.NDetecV
        };

        int N = ctx.Thetas.Count;
        var res = new double[2 * N];
        for (int i = 0; i < N; i++)
        {
            Projection.ProjectPoint(
                ctx.Thetas[i], ay, detX, detY, detZ,
                -R, offU, offV, gp,
                out double u, out double v);
            res[2 * i] = u - ctx.UObs[i];
            res[2 * i + 1] = v - ctx.VObs[i];
        }
        return res;
    }

    private static double[,] ComputeJacobian(double[] p, LmContext ctx)
    {
        int nObs = 2 * ctx.Thetas.Count;
        int np = ctx.NP;
        var J = new double[nObs, np];
        double eps = 1e-7;

        for (int j = 0; j < np; j++)
        {
            double[] pp = (double[])p.Clone();
            double[] pm = (double[])p.Clone();
            double step = Math.Max(eps, Math.Abs(p[j]) * eps);
            pp[j] = p[j] + step;
            pm[j] = p[j] - step;

            var rp = ComputeResiduals(pp, ctx);
            var rm = ComputeResiduals(pm, ctx);

            double denom = 2.0 * step;
            for (int i = 0; i < nObs; i++)
                J[i, j] = (rp[i] - rm[i]) / denom;
        }
        return J;
    }

    private static void ClampParams(double[] p, LmContext ctx)
    {
        for (int i = 0; i < BASE_PARAMS; i++)
            p[i] = Math.Clamp(p[i], LowerBase[i], UpperBase[i]);

        if (ctx.OptOffset)
        {
            p[ctx.IdxOffDU] = Math.Clamp(p[ctx.IdxOffDU], -50, 50);
            p[ctx.IdxOffDV] = Math.Clamp(p[ctx.IdxOffDV], -50, 50);
        }
        if (ctx.OptDist)
        {
            p[ctx.IdxDSO] = Math.Clamp(p[ctx.IdxDSO], 50, 1000);
            p[ctx.IdxDSD] = Math.Clamp(p[ctx.IdxDSD], 100, 2000);
            if (p[ctx.IdxDSD] <= p[ctx.IdxDSO])
                p[ctx.IdxDSD] = p[ctx.IdxDSO] + 10;
        }
    }

    private bool SolveLM(double[] p, LmContext ctx, int maxIter, double tol)
    {
        double lambda = 1e-3;
        var res = ComputeResiduals(p, ctx);
        int m = res.Length;
        int np = ctx.NP;
        double cost = res.Sum(r => r * r);

        for (int iter = 0; iter < maxIter; iter++)
        {
            var J = ComputeJacobian(p, ctx);

            // JtJ = J^T * J, Jtr = J^T * r
            var JtJ = new double[np, np];
            var Jtr = new double[np];
            for (int i = 0; i < np; i++)
            {
                for (int j = i; j < np; j++)
                {
                    double sum = 0;
                    for (int k = 0; k < m; k++)
                        sum += J[k, i] * J[k, j];
                    JtJ[i, j] = sum;
                    JtJ[j, i] = sum;
                }
                double s = 0;
                for (int k = 0; k < m; k++)
                    s += J[k, i] * res[k];
                Jtr[i] = s;
            }

            // A = JtJ + lambda * diag(1 + JtJ)
            var A = new double[np, np];
            var b = new double[np];
            Array.Copy(JtJ, A, JtJ.Length);
            for (int i = 0; i < np; i++)
            {
                A[i, i] += lambda * (1.0 + JtJ[i, i]);
                b[i] = -Jtr[i];
            }

            // 解线性方程组 (Cholesky)
            if (!SolveLinear(A, b, np, out var delta))
            {
                lambda *= 10;
                continue;
            }

            // 尝试更新
            var np2 = new double[np];
            for (int i = 0; i < np; i++)
                np2[i] = p[i] + delta[i];
            ClampParams(np2, ctx);

            var newRes = ComputeResiduals(np2, ctx);
            double newCost = newRes.Sum(r => r * r);

            if (newCost < cost)
            {
                double improvement = cost - newCost;
                Array.Copy(np2, p, np);
                res = newRes;
                cost = newCost;
                lambda *= 0.3;
                if (lambda < 1e-12) lambda = 1e-12;

                OnProgress?.Invoke(iter, Math.Sqrt(cost / m), p);

                if (improvement < tol)
                    return true;
            }
            else
            {
                lambda *= 10;
                if (lambda > 1e12) lambda = 1e12;
            }
        }
        return false;
    }

    /// <summary>
    /// 简单的 Cholesky 分解求解 Ax = b
    /// </summary>
    private static bool SolveLinear(double[,] A, double[] b, int n, out double[] x)
    {
        x = new double[n];
        var L = new double[n, n];

        // Cholesky: A = L * L^T
        for (int i = 0; i < n; i++)
        {
            for (int j = 0; j <= i; j++)
            {
                double sum = 0;
                for (int k = 0; k < j; k++)
                    sum += L[i, k] * L[j, k];

                if (i == j)
                {
                    double val = A[i, i] - sum;
                    if (val <= 0) return false;
                    L[i, j] = Math.Sqrt(val);
                }
                else
                {
                    L[i, j] = (A[i, j] - sum) / L[j, j];
                }
            }
        }

        // 前代: L * y = b
        var y = new double[n];
        for (int i = 0; i < n; i++)
        {
            double sum = 0;
            for (int k = 0; k < i; k++)
                sum += L[i, k] * y[k];
            y[i] = (b[i] - sum) / L[i, i];
        }

        // 回代: L^T * x = y
        for (int i = n - 1; i >= 0; i--)
        {
            double sum = 0;
            for (int k = i + 1; k < n; k++)
                sum += L[k, i] * x[k];
            x[i] = (y[i] - sum) / L[i, i];
        }

        return true;
    }

    #endregion
}