"""MultistateAnalyzer orchestrator."""

from typing import Any, Dict, List

import numpy as np
import pandas as pd

from ...config import PhenoClusterConfig
from ...utils.logging import get_logger
from .hazard_ratios import extract_hazard_ratios
from .pathways import analyze_pathway_frequencies
from .trajectory import TrajectoryBuilder
from .transition_hazards import TransitionHazardFitter
from .types import MonteCarloResults, MultistateResults


class MultistateAnalyzer:
    """Multi-state survival analysis with phenotype-specific transitions.

    Uses lifelines CoxPHFitter for transition-specific hazard modeling and
    Cox proportional hazards for hazard ratio estimation.
    """

    MAX_REASONABLE_HR = 1000

    def __init__(self, config: PhenoClusterConfig, n_clusters: int, reference_phenotype: int = 0):
        self.config = config
        self.n_clusters = n_clusters
        self.reference_phenotype = reference_phenotype
        self.logger = get_logger("multistate", config)
        self.multistate_config = config.multistate

        self._traj_builder = TrajectoryBuilder(config, n_clusters, reference_phenotype, self.logger)

    def prepare_trajectories(self, data, labels, baseline_confounders=None):
        """Delegate to TrajectoryBuilder."""
        return self._traj_builder.prepare_trajectories(data, labels, baseline_confounders)

    def fit_transition_model(self, trajectories):
        """Fit multistate transition-specific hazard model."""
        self.logger.info("Fitting transition-specific hazard model (lifelines)...")

        if not trajectories:
            raise ValueError("No trajectories provided for transition model fitting")

        terminal_states = self._traj_builder.get_terminal_state_ids()
        covariate_names = trajectories[0].covariates.index.tolist()

        model = TransitionHazardFitter(
            terminal_states=terminal_states,
            covariate_names=covariate_names,
            logger=self.logger,
            random_state=self.config.random_state,
        )
        model.fit(trajectories)

        n_fitted = sum(len(targets) for targets in model.state_models.values())
        self.logger.info(f"Fitted {n_fitted} transition-specific Cox models")
        return model

    def extract_hazard_ratios(self, trajectories):
        """Fit Cox PH models for each transition and extract hazard ratios."""
        return extract_hazard_ratios(
            self.multistate_config.transitions,
            trajectories,
            self.n_clusters,
            self.reference_phenotype,
            self.config.inference,
            self.MAX_REASONABLE_HR,
            self.logger,
            min_events_per_transition=(self.multistate_config.min_events_per_transition),
            baseline_confounders=(self.multistate_config.baseline_confounders),
        )

    def run_monte_carlo_predictions(
        self,
        model,
        trajectories,
        phenotype_indices,
        n_simulations=None,
        time_points=None,
        max_transitions=None,
    ) -> MonteCarloResults:
        """Estimate state occupation probabilities per phenotype using MC."""
        self.logger.info("Running Monte Carlo predictions...")

        times = time_points or self.multistate_config.monte_carlo_time_points
        max_trans = max_transitions or self.multistate_config.max_transitions_per_path

        # Validate time points against observed follow-up
        if trajectories:
            max_observed = max(
                t.time_at_each_state[-1] if t.time_at_each_state else 0 for t in trajectories
            )
            if max_observed > 0:
                beyond = [t for t in times if t > max_observed]
                if beyond:
                    self.logger.warning(
                        f"MC time points {beyond} exceed max observed "
                        f"follow-up ({max_observed:.1f}). "
                        f"Extrapolated probabilities may be unreliable."
                    )

        initial_state = self._traj_builder.get_initial_state().id
        all_state_ids = [s.id for s in self.multistate_config.states]
        n_sims = n_simulations or self.multistate_config.monte_carlo_n_simulations

        traj_by_id = {t.sample_id: t for t in trajectories}
        state_probs: Dict[int, Dict[int, List[float]]] = {}
        state_probs_lower: Dict[int, Dict[int, List[float]]] = {}
        state_probs_upper: Dict[int, Dict[int, List[float]]] = {}
        n_boot = 200

        for phenotype in range(self.n_clusters):
            state_probs[phenotype] = {s: [0.0] * len(times) for s in all_state_ids}
            state_probs_lower[phenotype] = {s: [0.0] * len(times) for s in all_state_ids}
            state_probs_upper[phenotype] = {s: [0.0] * len(times) for s in all_state_ids}

            patient_idxs = phenotype_indices.get(phenotype, [])
            if not patient_idxs:
                self.logger.warning(f"No patients in phenotype {phenotype}")
                continue

            patient_trajs = [traj_by_id[i] for i in patient_idxs if i in traj_by_id]
            if not patient_trajs:
                self.logger.warning(f"No valid trajectories for phenotype {phenotype}")
                continue

            # Cap per-patient simulations so total work stays roughly the same
            per_patient_sims = max(50, n_sims // len(patient_trajs))

            # Marginalise over individual patients
            patient_probs = []
            for traj in patient_trajs:
                covs = traj.covariates.values
                try:
                    probs = self._mc_for_covariates(
                        model,
                        covs,
                        initial_state,
                        all_state_ids,
                        times,
                        max_trans,
                        per_patient_sims,
                    )
                    patient_probs.append(probs)
                except Exception:
                    continue

            if not patient_probs:
                self.logger.warning(f"Monte Carlo failed for all patients in phenotype {phenotype}")
                continue

            # Average across patients (marginalisation)
            for s in all_state_ids:
                vals = [p[s] for p in patient_probs]
                state_probs[phenotype][s] = [
                    float(np.mean([v[t_idx] for v in vals])) for t_idx in range(len(times))
                ]

            # Bootstrap CIs
            rng = np.random.RandomState(self.config.random_state)
            boot_probs: Dict[int, List[List[float]]] = {s: [] for s in all_state_ids}
            for _ in range(n_boot):
                boot_idx = rng.choice(len(patient_probs), len(patient_probs), replace=True)
                for s in all_state_ids:
                    vals = [patient_probs[boot_idx[j]][s] for j in range(len(boot_idx))]
                    boot_probs[s].append(
                        [float(np.mean([v[t_idx] for v in vals])) for t_idx in range(len(times))]
                    )

            for s in all_state_ids:
                arr = np.array(boot_probs[s])  # shape (n_boot, n_times)
                state_probs_lower[phenotype][s] = np.percentile(arr, 2.5, axis=0).tolist()
                state_probs_upper[phenotype][s] = np.percentile(arr, 97.5, axis=0).tolist()

            self.logger.info(
                f"Phenotype {phenotype}: marginalised over {len(patient_probs)} patients "
                f"({per_patient_sims} sims each), {n_boot} bootstrap resamples"
            )

        self.logger.info(f"Monte Carlo completed: {len(times)} time points")
        return MonteCarloResults(
            time_points=times,
            state_probabilities=state_probs,
            n_simulations=n_sims,
            simulation_summary={
                "max_transitions": max_trans,
                "initial_state": initial_state,
                "n_phenotypes": self.n_clusters,
                "method": "marginalised_coxph",
                "n_bootstrap": n_boot,
            },
            state_probabilities_lower=state_probs_lower,
            state_probabilities_upper=state_probs_upper,
        )

    def _mc_for_covariates(
        self,
        model,
        sample_covariates,
        initial_state,
        all_state_ids,
        times,
        max_trans,
        n_sims,
    ) -> Dict[int, List[float]]:
        """Run MC simulation for a single covariate vector.

        Returns a dict mapping state_id -> [prob_at_each_time_point].
        """
        simulated = model.run_monte_carlo_simulation(
            sample_covariates=sample_covariates,
            origin_state=initial_state,
            current_time=0,
            n_random_samples=n_sims,
            max_transitions=max_trans,
            n_jobs=1,
        )
        result: Dict[int, List[float]] = {s: [0.0] * len(times) for s in all_state_ids}
        for t_idx, t in enumerate(times):
            state_counts = {s: 0 for s in all_state_ids}
            for sim_traj in simulated:
                st = sim_traj.state_at_time(t, initial_state)
                state_counts[st] = state_counts.get(st, 0) + 1
            for s in all_state_ids:
                result[s][t_idx] = state_counts.get(s, 0) / n_sims
        return result

    def analyze_pathway_frequencies(self, data, labels, trajectories=None):
        """Analyze pathway frequencies by phenotype."""
        return analyze_pathway_frequencies(
            data,
            labels,
            self.n_clusters,
            self._traj_builder._state_by_id,
            trajectories=trajectories,
            determine_state_sequence=self._traj_builder._determine_state_sequence,
            logger=self.logger,
        )

    def run_full_analysis(self, data: pd.DataFrame, labels: np.ndarray) -> MultistateResults:
        """Run complete multistate analysis pipeline."""
        self.logger.info("MULTISTATE ANALYSIS (lifelines)")

        missing_cols = self._check_columns(data)
        if missing_cols:
            self.logger.error(f"Missing required columns: {missing_cols}")
            return MultistateResults(
                model_summary={"error": f"Missing columns: {missing_cols}"},
                transition_results={},
                pathway_results=[],
            )

        trajectories, phenotype_indices = self.prepare_trajectories(data, labels)
        if len(trajectories) == 0:
            self.logger.error("No valid trajectories created")
            return MultistateResults(
                model_summary={"error": "No valid patient trajectories"},
                transition_results={},
                pathway_results=[],
            )

        try:
            model = self.fit_transition_model(trajectories)
        except Exception as e:
            self.logger.error(f"Transition model fitting failed: {e}")
            return MultistateResults(
                model_summary={"error": f"Model fitting failed: {e}"},
                transition_results={},
                pathway_results=[],
            )

        transition_results = self.extract_hazard_ratios(trajectories)

        mc_results = None
        if model.state_models:
            try:
                mc_results = self.run_monte_carlo_predictions(
                    model, trajectories, phenotype_indices
                )
            except Exception as e:
                self.logger.warning(f"Monte Carlo predictions failed: {e}")
        else:
            self.logger.warning("No fitted models available for Monte Carlo simulation")

        pathway_results = self.analyze_pathway_frequencies(data, labels, trajectories=trajectories)

        model_summary = {
            "n_patients": len(data),
            "n_trajectories": len(trajectories),
            "n_states": len(self.multistate_config.states),
            "n_allowed_transitions": len(self.multistate_config.transitions),
            "n_transitions_fitted": len(transition_results),
            "n_pathways_observed": len(pathway_results),
            "monte_carlo_n_simulations": self.multistate_config.monte_carlo_n_simulations,
            "transitions_fitted": list(transition_results.keys()),
            "analysis_method": "lifelines_cox_ph",
            "inference": "Cox proportional hazards",
        }

        self.logger.info("MULTISTATE ANALYSIS COMPLETE")
        self.logger.info(f"  Patients: {model_summary['n_patients']}")
        self.logger.info(f"  Transitions fitted: {len(transition_results)}")
        self.logger.info(f"  Unique pathways: {len(pathway_results)}")
        if mc_results:
            self.logger.info(f"  Monte Carlo simulations: {mc_results.n_simulations}")

        return MultistateResults(
            model_summary=model_summary,
            transition_results=transition_results,
            pathway_results=pathway_results,
            state_occupation_probabilities=mc_results,
            fitted_model=model,
        )

    def _check_columns(self, data):
        """Validate that required columns exist."""
        missing = []
        for state in self.multistate_config.states:
            if state.event_column and state.event_column not in data.columns:
                missing.append(state.event_column)
            if state.time_column and state.time_column not in data.columns:
                missing.append(state.time_column)
        return missing

    def results_to_dict(self, results: MultistateResults) -> Dict[str, Any]:
        """Convert MultistateResults to a JSON-serializable dictionary."""
        result_dict = {
            "model_summary": results.model_summary,
            "transition_results": {
                name: {
                    "transition_name": tr.transition_name,
                    "from_state": tr.from_state,
                    "to_state": tr.to_state,
                    "n_events": tr.n_events,
                    "n_at_risk": tr.n_at_risk,
                    "phenotype_effects": tr.phenotype_effects,
                    "covariate_effects": tr.covariate_effects,
                }
                for name, tr in results.transition_results.items()
            },
            "pathway_results": [
                {
                    "pathway": pr.pathway,
                    "state_names": pr.state_names,
                    "counts_by_phenotype": {str(k): v for k, v in pr.counts_by_phenotype.items()},
                    "total_count": pr.total_count,
                    "fraction_by_phenotype": {
                        str(k): v for k, v in pr.fraction_by_phenotype.items()
                    },
                }
                for pr in results.pathway_results
            ],
        }

        if results.state_occupation_probabilities:
            mc = results.state_occupation_probabilities
            mc_dict: Dict[str, Any] = {
                "time_points": mc.time_points,
                "n_simulations": mc.n_simulations,
                "by_phenotype": {
                    str(phenotype): {str(state): probs for state, probs in state_probs.items()}
                    for phenotype, state_probs in mc.state_probabilities.items()
                },
                "simulation_summary": mc.simulation_summary,
            }
            if mc.state_probabilities_lower is not None:
                mc_dict["by_phenotype_lower"] = {
                    str(phenotype): {str(state): probs for state, probs in state_probs.items()}
                    for phenotype, state_probs in mc.state_probabilities_lower.items()
                }
            if mc.state_probabilities_upper is not None:
                mc_dict["by_phenotype_upper"] = {
                    str(phenotype): {str(state): probs for state, probs in state_probs.items()}
                    for phenotype, state_probs in mc.state_probabilities_upper.items()
                }
            result_dict["state_occupation_probabilities"] = mc_dict
        else:
            result_dict["state_occupation_probabilities"] = None

        return result_dict
