"""
Lunaron Evolution Engine - Candidate Evaluator
Evaluates challenger models against champion baseline on Holdout Dataset.
Strict Regression Gate:
- If ANY stem degrades by > 5.0%, regression fails.
- Prevents "Vocal improved at the expense of Drums" regressions.
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any
import torch
import torch.nn.functional as F

@dataclass
class RegressionResult:
    passed: bool
    overall_improved: bool
    stem_regressions: List[str]
    reasons: List[str]
    summary_text: str

@dataclass
class EvaluationReport:
    champion_id: str
    challenger_id: str
    stem_metrics: Dict[str, Dict[str, float]]
    regression_result: RegressionResult
    overall_quality_delta_pct: float
    leakage_delta_pct: float
    mixture_consistency_delta_pct: float

class CandidateEvaluator:
    """
    Offline validation & regression test engine.
    Ensures safe, measured model advancement without degrading any stem.
    """

    @staticmethod
    def compute_stem_metrics(
        pred_stem: torch.Tensor,
        ref_stem: torch.Tensor,
        eps: float = 1e-8
    ) -> Dict[str, float]:
        """Calculates energy, SDR proxy, and spectral convergence for a stem."""
        min_len = min(pred_stem.shape[-1], ref_stem.shape[-1])
        p = pred_stem[..., :min_len]
        r = ref_stem[..., :min_len]

        p_mono = p.mean(dim=0) if p.dim() > 1 else p
        r_mono = r.mean(dim=0) if r.dim() > 1 else r

        l1_err = float(F.l1_loss(p_mono, r_mono).item())
        p_energy = float(torch.sum(p_mono ** 2).item())
        r_energy = float(torch.sum(r_mono ** 2).item())

        # Spectral difference
        win = torch.hann_window(2048, device=p.device)
        stft_p = torch.abs(torch.stft(p_mono, n_fft=2048, hop_length=512, window=win, return_complex=True)) + eps
        stft_r = torch.abs(torch.stft(r_mono, n_fft=2048, hop_length=512, window=win, return_complex=True)) + eps
        spectral_conv = float((torch.norm(stft_p - stft_r, p="fro") / (torch.norm(stft_r, p="fro") + eps)).item())

        return {
            "l1_error": round(l1_err, 5),
            "spectral_convergence": round(spectral_conv, 4),
            "energy_ratio": round(p_energy / (r_energy + eps), 4)
        }

    @classmethod
    def check_regression(
        cls,
        champion_metrics: Dict[str, Dict[str, float]],
        challenger_metrics: Dict[str, Dict[str, float]],
        tolerance_pct: float = 5.0
    ) -> RegressionResult:
        """
        Verifies whether challenger has any regression compared to champion.
        A stem is considered regressed if spectral convergence or error is worse by > tolerance_pct.
        """
        regressed_stems = []
        reasons = []

        total_champ_sc = 0.0
        total_chall_sc = 0.0
        count = 0

        for stem_name, c_m in champion_metrics.items():
            if stem_name not in challenger_metrics:
                continue
            h_m = challenger_metrics[stem_name]
            count += 1

            c_sc = c_m.get("spectral_convergence", 1.0)
            h_sc = h_m.get("spectral_convergence", 1.0)
            total_champ_sc += c_sc
            total_chall_sc += h_sc

            # Higher spectral convergence means more error
            degradation = ((h_sc - c_sc) / max(c_sc, 1e-6)) * 100.0
            if degradation > tolerance_pct:
                regressed_stems.append(stem_name)
                reasons.append(f"Stem '{stem_name}' degraded by {round(degradation, 1)}% in spectral error.")

        overall_improved = total_chall_sc < total_champ_sc if count > 0 else False
        passed = (len(regressed_stems) == 0) and overall_improved

        if not overall_improved:
            reasons.append("Challenger overall error is not lower than Champion.")

        summary = (
            f"Regression Status: {'PASS' if passed else 'FAIL'}. "
            f"Active Regressions: {len(regressed_stems)} stems."
        )

        return RegressionResult(
            passed=passed,
            overall_improved=overall_improved,
            stem_regressions=regressed_stems,
            reasons=reasons,
            summary_text=summary
        )

    @classmethod
    def generate_report_text(cls, report: EvaluationReport) -> str:
        """Generates formatted text report following the design spec."""
        reg = report.regression_result
        return f"""
========================================
Lunaron Evolution Engine Evaluation Report
========================================
Previous Champion: {report.champion_id}
Candidate        : {report.challenger_id}

Overall Quality Delta    : {report.overall_quality_delta_pct:+.1f}%
Leakage Delta            : {report.leakage_delta_pct:+.1f}%
Mixture Consistency Delta: {report.mixture_consistency_delta_pct:+.1f}%

Regression Gate Status   : {'PASS' if reg.passed else 'FAIL'}
Stem Regressions Detected: {', '.join(reg.stem_regressions) if reg.stem_regressions else 'None'}
{chr(10).join(['- ' + r for r in reg.reasons])}
========================================
"""