from __future__ import annotations

import logging
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass

import pandas as pd

from signalab.analytics.report import Report
from signalab.backtest_engine import BacktestEngine
from signalab.data.market_data import MarketData
from signalab.experiment.param_grid import ExperimentConfig, ParamGrid

logger = logging.getLogger(__name__)


def _run_single_experiment(
    run_id: int,
    config: ExperimentConfig,
    data: MarketData,
    engine_verbose: bool,
    engine_progress_log_interval: pd.Timedelta | None,
    engine_results_path: str | None,
    from_date: pd.Timestamp | None,
    to_date: pd.Timestamp | None,
) -> ExperimentRunResult:
    """Execute one experiment configuration.

    Parameters
    ----------
    run_id : int
        Sequential run identifier.
    config : ExperimentConfig
        Run configuration.
    data : MarketData
        Market data input.
    engine_verbose : bool
        Engine verbosity flag.
    engine_progress_log_interval : pd.Timedelta | None
        Engine progress log interval.
    engine_results_path : str | None
        Optional results folder path.
    from_date : pd.Timestamp | None
        Optional run start timestamp.
    to_date : pd.Timestamp | None
        Optional run end timestamp.

    Returns
    -------
    ExperimentRunResult
        Executed run result.
    """
    engine = BacktestEngine(
        verbose=engine_verbose,
        progress_log_interval=engine_progress_log_interval,
        results_path=engine_results_path,
    )
    report = engine.run(
        data=data,
        indicators=config.make_indicators(),
        strategies=[config.make_strategy()],
        accounts=config.make_account(),
        individual_accounts=config.individual_accounts,
        from_date=from_date,
        to_date=to_date,
    )

    metrics_by_account = {
        account_label: account_report.metrics()
        for account_label, account_report in report.account_reports.items()
    }

    # Drop non-pickleable runtime row objects (pandas itertuples namedtuple class)
    # before sending results back from process workers.
    for account_report in report.account_reports.values():
        account_report.account._current_market = None

    return ExperimentRunResult(
        run_id=run_id,
        config=config,
        report=report,
        metrics_by_account=metrics_by_account,
    )


@dataclass(frozen=True)
class ExperimentRunResult:
    """Container for a single executed experiment run."""

    run_id: int
    config: ExperimentConfig
    report: Report
    metrics_by_account: dict[str, dict[str, float]]


class ExperimentResult:
    """Collection of executed experiment runs."""

    def __init__(self, runs: list[ExperimentRunResult]):
        """Initialize experiment result collection.

        Parameters
        ----------
        runs : list[ExperimentRunResult]
            Executed run results.
        """
        self.runs = runs

    def __len__(self) -> int:
        """Return number of runs.

        Returns
        -------
        int
            Run count.
        """
        return len(self.runs)

    def summary(self) -> pd.DataFrame:
        """Build summary dataframe from all runs.

        Returns
        -------
        pd.DataFrame
            Per-run (and per-account) summary table.
        """
        rows: list[dict[str, object]] = []
        for run in self.runs:
            if run.metrics_by_account:
                for account_label, metrics in run.metrics_by_account.items():
                    row: dict[str, object] = {
                        "run_id": run.run_id,
                        "strategy_cls": run.config.strategy_cls.__name__,
                        "strategy_params": run.config.strategy_params,
                        "account_label": account_label,
                        "account_params": run.config.account_params,
                        "individual_accounts": run.config.individual_accounts,
                        "indicators": [
                            {
                                "cls": indicator.cls.__name__,
                                "params": indicator.params,
                            }
                            for indicator in run.config.indicators
                        ],
                    }
                    row.update(metrics)
                    rows.append(row)
            else:
                rows.append(
                    {
                        "run_id": run.run_id,
                        "strategy_cls": run.config.strategy_cls.__name__,
                        "strategy_params": run.config.strategy_params,
                        "account_label": None,
                        "account_params": run.config.account_params,
                        "individual_accounts": run.config.individual_accounts,
                        "indicators": [
                            {
                                "cls": indicator.cls.__name__,
                                "params": indicator.params,
                            }
                            for indicator in run.config.indicators
                        ],
                    }
                )
        return pd.DataFrame(rows)


class ExperimentRunner:
    """Runner for executing parameter-grid experiments."""

    def __init__(self, data: MarketData, backtest_engine: BacktestEngine):
        """Initialize experiment runner.

        Parameters
        ----------
        data : MarketData
            Market dataset used in all runs.
        backtest_engine : BacktestEngine
            Engine instance used as template for run execution.
        """
        self._data = data
        self._backtest_engine = backtest_engine

    def run(
        self,
        param_grid: ParamGrid,
        from_date: pd.Timestamp | None = None,
        to_date: pd.Timestamp | None = None,
        n_jobs: int = 0,
    ) -> ExperimentResult:
        """Run all configurations from a parameter grid.

        Parameters
        ----------
        param_grid : ParamGrid
            Built grid of experiment configurations.
        from_date : pd.Timestamp | None, default=None
            Optional inclusive start timestamp.
        to_date : pd.Timestamp | None, default=None
            Optional inclusive end timestamp.
        n_jobs : int, default=0
            Number of parallel worker processes. ``0`` runs sequentially.

        Returns
        -------
        ExperimentResult
            Executed run collection.
        """
        if n_jobs < 0:
            raise ValueError("n_jobs must be >= 0")

        if n_jobs == 0:
            return self._run_sequential(
                param_grid=param_grid,
                from_date=from_date,
                to_date=to_date,
            )

        return self._run_parallel(
            param_grid=param_grid,
            from_date=from_date,
            to_date=to_date,
            n_jobs=n_jobs,
        )

    def _run_sequential(
        self,
        param_grid: ParamGrid,
        from_date: pd.Timestamp | None,
        to_date: pd.Timestamp | None,
    ) -> ExperimentResult:
        """Run parameter grid sequentially.

        Returns
        -------
        ExperimentResult
            Executed run collection.
        """
        runs: list[ExperimentRunResult] = []

        for run_id, config in enumerate(param_grid):
            logger.info(
                "Running experiment configuration %d/%d", run_id + 1, len(param_grid)
            )
            report = self._backtest_engine.run(
                data=self._data,
                indicators=config.make_indicators(),
                strategies=[config.make_strategy()],
                accounts=config.make_account(),
                individual_accounts=config.individual_accounts,
                from_date=from_date,
                to_date=to_date,
            )

            metrics_by_account = {
                account_label: account_report.metrics()
                for account_label, account_report in report.account_reports.items()
            }

            runs.append(
                ExperimentRunResult(
                    run_id=run_id,
                    config=config,
                    report=report,
                    metrics_by_account=metrics_by_account,
                )
            )

        return ExperimentResult(runs)

    def _run_parallel(
        self,
        param_grid: ParamGrid,
        from_date: pd.Timestamp | None,
        to_date: pd.Timestamp | None,
        n_jobs: int,
    ) -> ExperimentResult:
        """Run parameter grid in parallel using process workers.

        Returns
        -------
        ExperimentResult
            Executed run collection.
        """
        engine_results_path: str | None = (
            str(self._backtest_engine.results_path)
            if self._backtest_engine.results_path is not None
            else None
        )

        runs_by_id: dict[int, ExperimentRunResult] = {}
        with ProcessPoolExecutor(max_workers=n_jobs) as executor:
            futures = [
                executor.submit(
                    _run_single_experiment,
                    run_id,
                    config,
                    self._data,
                    self._backtest_engine.verbose,
                    self._backtest_engine.progress_log_interval,
                    engine_results_path,
                    from_date,
                    to_date,
                )
                for run_id, config in enumerate(param_grid)
            ]

            for future in as_completed(futures):
                run = future.result()
                runs_by_id[run.run_id] = run

        ordered_runs = [runs_by_id[i] for i in sorted(runs_by_id)]
        return ExperimentResult(ordered_runs)
