"""Step 8: Perform statistical analysis (ANOVA and heritability).

NOTE: This step only generates data files (CSVs, JSONs).
Heritability plots are generated by the generate_static_figures step to avoid duplication.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Optional

import pandas as pd

from sleap_roots_analyze.pipeline.core import BaseStep, StepResult
from sleap_roots_analyze.statistics import (
    calculate_heritability_estimates,
    calculate_trait_statistics,
    perform_anova_by_genotype,
)


class StatisticalAnalysisStep(BaseStep):
    """Perform statistical analysis on cleaned data.

    This step calculates:
    - Basic trait statistics (mean, std, min, max, etc.)
    - ANOVA results for each trait
    - Broad-sense heritability (H²) for each trait (when enabled)

    Heritability calculation is controlled by
    ``config.statistics.calculate_heritability``. When the flag is ``False``,
    heritability is skipped entirely and no heritability CSV is produced.

    Outputs:
        - 08_trait_statistics.json: Basic statistics for each trait
        - 08_anova_results.csv: ANOVA F-statistics and p-values
        - 08_heritability_results.csv: Heritability estimates and significance
          (only generated when ``config.statistics.calculate_heritability``
          is enabled)
        - 08_statistical_analysis_summary.json: Combined summary
    """

    def __init__(self):
        """Initialize StatisticalAnalysisStep."""
        super().__init__(
            step_name="StatisticalAnalysis",
            description="Perform statistical analysis (ANOVA and heritability)",
        )

    def execute(
        self,
        data: Any,
        config: Any,
        run_dir: Path,
        prev_result: Optional[StepResult] = None,
    ) -> StepResult:
        """Execute statistical analysis.

        Args:
            data: DataFrame from previous step (RemoveOutliersStep).
            config: Pipeline configuration with:
                - columns.genotype: Genotype column name
                - columns.replicate: Replicate column name (optional)
            run_dir: Directory to save outputs.
            prev_result: Result from previous step.

        Returns:
            StepResult with original DataFrame and statistical results in metadata.
        """
        df = data.copy()

        # Get valid trait columns from previous step
        trait_cols = prev_result.metadata.get("valid_trait_names")
        if trait_cols is None:
            # Try to get from earlier step
            trait_cols = prev_result.metadata.get("trait_names")

        if trait_cols is None:
            raise ValueError(
                "Could not find trait column names in previous step metadata. "
                "Expected 'valid_trait_names' or 'trait_names'."
            )

        # After CleanupTraitsStep (Step 02), column names are sanitized to standard names
        # Use the sanitized names consistently across all subsequent steps
        genotype_col = "Genotype"
        replicate_col = "Replicate"

        # 1. Calculate basic trait statistics
        trait_stats = calculate_trait_statistics(df=df, trait_cols=trait_cols)

        # 2. Perform ANOVA for each trait
        anova_results = perform_anova_by_genotype(
            df=df, trait_cols=trait_cols, genotype_col=genotype_col, alpha=0.05
        )

        # Convert ANOVA results to DataFrame for saving
        anova_records = []
        for trait, result in anova_results.items():
            if trait == "__calculation_metadata__":
                continue

            # CRITICAL: Check if result is a string (error message) FIRST
            # before checking for "error" key, to avoid AttributeError
            if isinstance(result, str):
                # Error case: result is an error message string
                anova_records.append(
                    {
                        "trait": trait,
                        "f_statistic": None,
                        "p_value": None,
                        "eta_squared": None,
                        "significant": None,
                        "n_groups": None,
                        "total_n": None,
                        "error": result,
                    }
                )
            elif "error" in result:
                # Error case: result is a dict with "error" key
                anova_records.append(
                    {
                        "trait": trait,
                        "f_statistic": None,
                        "p_value": None,
                        "eta_squared": None,
                        "significant": None,
                        "n_groups": None,
                        "total_n": None,
                        "error": result["error"],
                    }
                )
            else:
                # Success case: result is a dict with statistics
                anova_records.append(
                    {
                        "trait": trait,
                        "f_statistic": result.get("f_statistic"),
                        "p_value": result.get("p_value"),
                        "eta_squared": result.get("eta_squared"),
                        "significant": result.get("significant"),
                        "n_groups": result.get("n_groups"),
                        "total_n": result.get("total_n"),
                        "error": None,
                    }
                )

        anova_df = pd.DataFrame(anova_records)

        # 3. Calculate heritability for each trait (if enabled)
        statistics_config = getattr(config, "statistics", None)
        calculate_heritability = (
            statistics_config.calculate_heritability
            if statistics_config is not None
            else True
        )

        if calculate_heritability:
            heritability_results = calculate_heritability_estimates(
                df=df,
                trait_cols=trait_cols,
                genotype_col=genotype_col,
                replicate_col=replicate_col,
                force_method=None,  # Use default mixed model approach
                remove_low_h2=False,  # Don't remove yet, that's Step 9
            )

            # Convert heritability results to DataFrame for saving
            heritability_records = []
            for trait, result in heritability_results.items():
                if trait == "__calculation_metadata__":
                    continue
                if "error" in result:
                    heritability_records.append(
                        {
                            "trait": trait,
                            "heritability": None,
                            "var_genetic": None,
                            "var_residual": None,
                            "mean_n_reps": None,
                            "n_genotypes": None,
                            "n_observations": None,
                            "model_type": None,
                            "error": result["error"],
                        }
                    )
                else:
                    heritability_records.append(
                        {
                            "trait": trait,
                            "heritability": result.get("heritability"),
                            "var_genetic": result.get("var_genetic"),
                            "var_residual": result.get("var_residual"),
                            "mean_n_reps": result.get("mean_n_reps"),
                            "n_genotypes": result.get("n_genotypes"),
                            "n_observations": result.get("n_observations"),
                            "model_type": result.get("model_type"),
                            "error": None,
                        }
                    )

            heritability_df = pd.DataFrame(heritability_records)
        else:
            # Heritability disabled: empty results, no CSV
            heritability_results = {}
            heritability_df = pd.DataFrame()

        # 4. Create combined summary
        # Count valid ANOVA results (filter out strings and error dicts)
        valid_anova = [
            r
            for t, r in anova_results.items()
            if t != "__calculation_metadata__"
            and not isinstance(r, str)
            and "error" not in r
        ]
        significant_p005 = sum(
            1
            for r in valid_anova
            if r.get("p_value") is not None and r["p_value"] < 0.05
        )
        significant_p001 = sum(
            1
            for r in valid_anova
            if r.get("p_value") is not None and r["p_value"] < 0.01
        )

        if calculate_heritability:
            # Count valid heritability results
            valid_h2 = [
                r
                for t, r in heritability_results.items()
                if t != "__calculation_metadata__"
                and not isinstance(r, str)
                and "error" not in r
            ]
            high_h2_60 = sum(
                1
                for r in valid_h2
                if r.get("heritability") is not None and r["heritability"] >= 0.60
            )
            mean_h2 = (
                sum(
                    r["heritability"]
                    for r in valid_h2
                    if r.get("heritability") is not None
                )
                / len(valid_h2)
                if valid_h2
                else None
            )
            heritability_summary = {
                "traits_with_results": len(valid_h2),
                "high_heritability_traits_h60": high_h2_60,
                "mean_heritability": mean_h2,
            }
        else:
            heritability_summary = {"skipped": True}

        summary = {
            "n_traits_analyzed": len(trait_cols),
            "n_samples": len(df),
            "n_genotypes": df[genotype_col].nunique(),
            "traits_analyzed": trait_cols,
            "anova_summary": {
                "traits_with_results": len(valid_anova),
                "significant_traits_p005": significant_p005,
                "significant_traits_p001": significant_p001,
            },
            "heritability_summary": heritability_summary,
        }

        # Save outputs (to data/ directory per VIZ-OUTPUT-001)
        data_dir = run_dir / "data"
        data_dir.mkdir(parents=True, exist_ok=True)

        files = []
        files.append(self.save_json(trait_stats, "08_trait_statistics.json", data_dir))
        files.append(self.save_dataframe(anova_df, "08_anova_results.csv", data_dir))

        if calculate_heritability:
            files.append(
                self.save_dataframe(
                    heritability_df, "08_heritability_results.csv", data_dir
                )
            )

        files.append(
            self.save_json(summary, "08_statistical_analysis_summary.json", data_dir)
        )

        # NOTE: Heritability plots are generated by generate_static_figures step
        # to avoid duplication. This step only generates data files.

        # Create metadata, preserving previous step's metadata (e.g., image_paths)
        # This ensures metadata like image_paths flows through the pipeline
        metadata = {
            **prev_result.metadata,  # Preserve metadata from previous steps
            "trait_statistics": trait_stats,
            "anova_results": anova_results,
            "heritability_results": heritability_results,
            "summary": summary,
            "trait_names": trait_cols,  # Primary key (standardized)
            "valid_trait_names": trait_cols,  # For consistency
        }

        return StepResult(data=df, metadata=metadata, files_generated=files)
