using CeresSharp;

namespace CeresSharp.Test;

[TestFixture]
public class LossFunctionTests
{
    [Test]
    public void TrivialLoss_ShouldCreate()
    {
        using var loss = new TrivialLoss();
        Assert.That(loss, Is.Not.Null);
    }

    [Test]
    public void HuberLoss_ShouldCreate()
    {
        using var loss = new HuberLoss(1.0);
        Assert.That(loss, Is.Not.Null);
    }

    [Test]
    public void HuberLoss_WithDifferentScaling_ShouldCreate()
    {
        using var loss1 = new HuberLoss(0.5);
        using var loss2 = new HuberLoss(2.0);
        Assert.That(loss1, Is.Not.Null);
        Assert.That(loss2, Is.Not.Null);
    }

    [Test]
    public void CauchyLoss_ShouldCreate()
    {
        using var loss = new CauchyLoss(1.0);
        Assert.That(loss, Is.Not.Null);
    }

    [Test]
    public void SoftLOneLoss_ShouldCreate()
    {
        using var loss = new SoftLOneLoss(1.0);
        Assert.That(loss, Is.Not.Null);
    }

    [Test]
    public void ArctanLoss_ShouldCreate()
    {
        using var loss = new ArctanLoss(1.0);
        Assert.That(loss, Is.Not.Null);
    }

    [Test]
    public void TolerantLoss_ShouldCreate()
    {
        using var loss = new TolerantLoss(1.0, 2.0);
        Assert.That(loss, Is.Not.Null);
    }

    [Test]
    public void LossFunction_WithProblem_ShouldWork()
    {
        using var problem = new Problem();
        var parameters = new double[] { 1.0 };
        problem.AddParameterBlock(parameters, parameters.Length);
        
        var costFunction = new AutoDiffCostFunction(
            (parameters, residuals) =>
            {
                residuals[0] = parameters[0][0] - 1.0;
                return true;
            },
            numResiduals: 1,
            parameterBlockSizes: new[] { 1 });
        
        using var loss = new HuberLoss(1.0);
        var residualBlockId = problem.AddResidualBlock(
            costFunction,
            loss,
            parameterBlocks: new[] { parameters });
        
        Assert.That(residualBlockId, Is.Not.EqualTo(IntPtr.Zero), "Residual block ID should not be zero");
    }

    [Test]
    public void MultipleLossFunctions_ShouldWork()
    {
        using var problem = new Problem();
        var params1 = new double[] { 1.0 };
        var params2 = new double[] { 2.0 };
        problem.AddParameterBlock(params1, params1.Length);
        problem.AddParameterBlock(params2, params2.Length);
        
        var costFunction1 = new AutoDiffCostFunction(
            (parameters, residuals) =>
            {
                residuals[0] = parameters[0][0] - 1.0;
                return true;
            },
            numResiduals: 1,
            parameterBlockSizes: new[] { 1 });
        
        var costFunction2 = new AutoDiffCostFunction(
            (parameters, residuals) =>
            {
                residuals[0] = parameters[0][0] - 2.0;
                return true;
            },
            numResiduals: 1,
            parameterBlockSizes: new[] { 1 });
        
        using var loss1 = new HuberLoss(1.0);
        using var loss2 = new CauchyLoss(1.0);
        
        problem.AddResidualBlock(costFunction1, loss1, new[] { params1 });
        problem.AddResidualBlock(costFunction2, loss2, new[] { params2 });
        
        Assert.That(problem.NumResidualBlocks, Is.EqualTo(2));
    }
}