"""
Data comparison utilities for validating Python vs MATLAB outputs.

Provides statistical and numerical comparison functions to ensure
the Python implementation matches MATLAB results within acceptable tolerances.
"""

import numpy as np
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass
from enum import Enum


class ComparisonStatus(Enum):
    """Status of comparison between Python and MATLAB data."""
    IDENTICAL = "identical"
    EQUIVALENT = "equivalent"  # Within tolerance
    DIFFERENT = "different"
    MISSING_MATLAB = "missing_matlab"
    MISSING_PYTHON = "missing_python"
    ERROR = "error"


@dataclass
class ComparisonResult:
    """Result of comparing Python and MATLAB data."""
    status: ComparisonStatus
    field_name: str
    max_abs_diff: Optional[float] = None
    max_rel_diff: Optional[float] = None
    mean_abs_diff: Optional[float] = None
    rmse: Optional[float] = None
    correlation: Optional[float] = None
    matlab_shape: Optional[Tuple] = None
    python_shape: Optional[Tuple] = None
    matlab_range: Optional[Tuple[float, float]] = None
    python_range: Optional[Tuple[float, float]] = None
    message: str = ""

    def __str__(self) -> str:
        """Human-readable representation."""
        if self.status == ComparisonStatus.IDENTICAL:
            return f"✓ {self.field_name}: IDENTICAL"
        elif self.status == ComparisonStatus.EQUIVALENT:
            msg = f"✓ {self.field_name}: EQUIVALENT (within tolerance)\n"
            msg += f"  Max abs diff: {self.max_abs_diff:.2e}\n"
            msg += f"  Max rel diff: {self.max_rel_diff:.2%}\n"
            msg += f"  RMSE: {self.rmse:.2e}\n"
            msg += f"  Correlation: {self.correlation:.6f}"
            return msg
        elif self.status == ComparisonStatus.DIFFERENT:
            msg = f"✗ {self.field_name}: DIFFERENT (exceeds tolerance)\n"
            msg += f"  Max abs diff: {self.max_abs_diff:.2e}\n"
            msg += f"  Max rel diff: {self.max_rel_diff:.2%}\n"
            msg += f"  RMSE: {self.rmse:.2e}\n"
            msg += f"  MATLAB range: [{self.matlab_range[0]:.2e}, {self.matlab_range[1]:.2e}]\n"
            msg += f"  Python range: [{self.python_range[0]:.2e}, {self.python_range[1]:.2e}]"
            return msg
        else:
            return f"? {self.field_name}: {self.status.value} - {self.message}"


class DataComparator:
    """Compare Python and MATLAB numerical data."""

    def __init__(self,
                 abs_tol: float = 1e-6,
                 rel_tol: float = 1e-4,
                 max_rel_tol: float = 0.01):  # 1%
        """
        Initialize comparator with tolerance thresholds.

        Args:
            abs_tol: Absolute tolerance for differences
            rel_tol: Relative tolerance for differences (fraction)
            max_rel_tol: Maximum acceptable relative difference
        """
        self.abs_tol = abs_tol
        self.rel_tol = rel_tol
        self.max_rel_tol = max_rel_tol

    def compare_arrays(self,
                       matlab_data: np.ndarray,
                       python_data: np.ndarray,
                       field_name: str = "data") -> ComparisonResult:
        """
        Compare two numpy arrays (MATLAB vs Python).

        Args:
            matlab_data: Data from MATLAB processing
            python_data: Data from Python processing
            field_name: Name of the field being compared

        Returns:
            ComparisonResult with comparison statistics
        """
        # Check shapes
        if matlab_data.shape != python_data.shape:
            return ComparisonResult(
                status=ComparisonStatus.DIFFERENT,
                field_name=field_name,
                matlab_shape=matlab_data.shape,
                python_shape=python_data.shape,
                message=f"Shape mismatch: MATLAB {matlab_data.shape} vs Python {python_data.shape}"
            )

        # Check for NaN/Inf
        matlab_valid = np.isfinite(matlab_data)
        python_valid = np.isfinite(python_data)

        if not np.array_equal(matlab_valid, python_valid):
            return ComparisonResult(
                status=ComparisonStatus.DIFFERENT,
                field_name=field_name,
                message="NaN/Inf pattern mismatch between MATLAB and Python"
            )

        # Use only valid values for comparison
        valid_mask = matlab_valid & python_valid
        if not valid_mask.any():
            return ComparisonResult(
                status=ComparisonStatus.ERROR,
                field_name=field_name,
                message="No valid values to compare"
            )

        matlab_valid_data = matlab_data[valid_mask]
        python_valid_data = python_data[valid_mask]

        # Check if identical
        if np.array_equal(matlab_valid_data, python_valid_data):
            return ComparisonResult(
                status=ComparisonStatus.IDENTICAL,
                field_name=field_name,
                max_abs_diff=0.0,
                max_rel_diff=0.0,
                mean_abs_diff=0.0,
                rmse=0.0,
                correlation=1.0
            )

        # Calculate differences
        abs_diff = np.abs(matlab_valid_data - python_valid_data)
        max_abs_diff = np.max(abs_diff)
        mean_abs_diff = np.mean(abs_diff)

        # Calculate relative differences (avoid division by zero)
        matlab_abs = np.abs(matlab_valid_data)
        rel_diff = np.zeros_like(abs_diff)
        nonzero_mask = matlab_abs > self.abs_tol
        rel_diff[nonzero_mask] = abs_diff[nonzero_mask] / matlab_abs[nonzero_mask]
        max_rel_diff = np.max(rel_diff) if nonzero_mask.any() else 0.0

        # Calculate RMSE
        rmse = np.sqrt(np.mean((matlab_valid_data - python_valid_data) ** 2))

        # Calculate correlation
        if matlab_valid_data.std() > 0 and python_valid_data.std() > 0:
            correlation = np.corrcoef(matlab_valid_data.flatten(),
                                     python_valid_data.flatten())[0, 1]
        else:
            correlation = 1.0 if max_abs_diff < self.abs_tol else 0.0

        # Determine status
        if max_abs_diff < self.abs_tol and max_rel_diff < self.rel_tol:
            status = ComparisonStatus.EQUIVALENT
        elif max_rel_diff < self.max_rel_tol:
            status = ComparisonStatus.EQUIVALENT
        else:
            status = ComparisonStatus.DIFFERENT

        return ComparisonResult(
            status=status,
            field_name=field_name,
            max_abs_diff=max_abs_diff,
            max_rel_diff=max_rel_diff,
            mean_abs_diff=mean_abs_diff,
            rmse=rmse,
            correlation=correlation,
            matlab_shape=matlab_data.shape,
            python_shape=python_data.shape,
            matlab_range=(np.min(matlab_valid_data), np.max(matlab_valid_data)),
            python_range=(np.min(python_valid_data), np.max(python_valid_data))
        )

    def compare_scalars(self,
                       matlab_value: float,
                       python_value: float,
                       field_name: str = "value") -> ComparisonResult:
        """
        Compare two scalar values.

        Args:
            matlab_value: Scalar from MATLAB
            python_value: Scalar from Python
            field_name: Name of the field

        Returns:
            ComparisonResult
        """
        # Convert to arrays and compare
        matlab_arr = np.array([matlab_value])
        python_arr = np.array([python_value])
        return self.compare_arrays(matlab_arr, python_arr, field_name)

    def compare_records(self,
                       matlab_records: List[Dict[str, Any]],
                       python_records: List[Dict[str, Any]],
                       key_fields: List[str],
                       value_fields: List[str]) -> List[ComparisonResult]:
        """
        Compare lists of database records.

        Args:
            matlab_records: Records from MATLAB processing
            python_records: Records from Python processing
            key_fields: Fields to use for matching records (e.g., timestamp, node_id)
            value_fields: Fields to compare numerically

        Returns:
            List of ComparisonResult objects
        """
        results = []

        # Check record counts
        if len(matlab_records) != len(python_records):
            results.append(ComparisonResult(
                status=ComparisonStatus.DIFFERENT,
                field_name="record_count",
                message=f"Record count mismatch: MATLAB {len(matlab_records)} vs Python {len(python_records)}"
            ))
            return results

        # Match records by key fields
        matlab_dict = {}
        for record in matlab_records:
            key = tuple(record[f] for f in key_fields)
            matlab_dict[key] = record

        python_dict = {}
        for record in python_records:
            key = tuple(record[f] for f in key_fields)
            python_dict[key] = record

        # Find unmatched keys
        matlab_keys = set(matlab_dict.keys())
        python_keys = set(python_dict.keys())

        missing_in_python = matlab_keys - python_keys
        missing_in_matlab = python_keys - matlab_keys

        if missing_in_python:
            results.append(ComparisonResult(
                status=ComparisonStatus.MISSING_PYTHON,
                field_name="records",
                message=f"Missing {len(missing_in_python)} records in Python output"
            ))

        if missing_in_matlab:
            results.append(ComparisonResult(
                status=ComparisonStatus.MISSING_MATLAB,
                field_name="records",
                message=f"Missing {len(missing_in_matlab)} records in MATLAB output"
            ))

        # Compare matching records
        common_keys = matlab_keys & python_keys

        for field in value_fields:
            matlab_values = []
            python_values = []

            for key in sorted(common_keys):
                matlab_val = matlab_dict[key].get(field)
                python_val = python_dict[key].get(field)

                if matlab_val is not None and python_val is not None:
                    matlab_values.append(matlab_val)
                    python_values.append(python_val)

            if matlab_values and python_values:
                matlab_arr = np.array(matlab_values)
                python_arr = np.array(python_values)
                results.append(self.compare_arrays(matlab_arr, python_arr, field))

        return results


class ValidationReport:
    """Generate validation reports."""

    def __init__(self):
        self.results: List[ComparisonResult] = []

    def add_result(self, result: ComparisonResult):
        """Add a comparison result."""
        self.results.append(result)

    def add_results(self, results: List[ComparisonResult]):
        """Add multiple comparison results."""
        self.results.extend(results)

    def get_summary(self) -> Dict[str, int]:
        """Get summary counts by status."""
        summary = {status.value: 0 for status in ComparisonStatus}
        for result in self.results:
            summary[result.status.value] += 1
        return summary

    def is_valid(self) -> bool:
        """Check if validation passed (all identical or equivalent)."""
        for result in self.results:
            if result.status not in [ComparisonStatus.IDENTICAL,
                                    ComparisonStatus.EQUIVALENT]:
                return False
        return True

    def generate_report(self, include_equivalent: bool = False) -> str:
        """
        Generate human-readable report.

        Args:
            include_equivalent: Include details for equivalent (passing) comparisons

        Returns:
            Formatted report string
        """
        lines = []
        lines.append("=" * 80)
        lines.append("VALIDATION REPORT: Python vs MATLAB Output Comparison")
        lines.append("=" * 80)
        lines.append("")

        summary = self.get_summary()
        lines.append("SUMMARY:")
        lines.append(f"  ✓ Identical:  {summary['identical']}")
        lines.append(f"  ✓ Equivalent: {summary['equivalent']}")
        lines.append(f"  ✗ Different:  {summary['different']}")
        lines.append(f"  ? Missing (MATLAB): {summary['missing_matlab']}")
        lines.append(f"  ? Missing (Python): {summary['missing_python']}")
        lines.append(f"  ! Errors: {summary['error']}")
        lines.append("")

        if self.is_valid():
            lines.append("✓✓✓ VALIDATION PASSED ✓✓✓")
        else:
            lines.append("✗✗✗ VALIDATION FAILED ✗✗✗")
        lines.append("")

        # Detailed results
        lines.append("-" * 80)
        lines.append("DETAILED RESULTS:")
        lines.append("-" * 80)
        lines.append("")

        for result in self.results:
            if not include_equivalent and result.status in [ComparisonStatus.IDENTICAL,
                                                            ComparisonStatus.EQUIVALENT]:
                continue
            lines.append(str(result))
            lines.append("")

        return "\n".join(lines)

    def save_report(self, filepath: str, include_equivalent: bool = False):
        """Save report to file."""
        report = self.generate_report(include_equivalent)
        with open(filepath, 'w') as f:
            f.write(report)
