# -- Import
from winnow.calibration.calibration_features import (
    PrositFeatures,
    MassErrorFeature,
    RetentionTimeFeature,
    ChimericFeatures,
    BeamFeatures,
)
from winnow.calibration.calibrator import ProbabilityCalibrator
from winnow.datasets.calibration_dataset import CalibrationDataset
from winnow.datasets.data_loaders import (
    InstaNovoDatasetLoader,
    MZTabDatasetLoader,
    PointNovoDatasetLoader,
    WinnowDatasetLoader,
)
from winnow.fdr.nonparametric import NonParametricFDRControl
from winnow.fdr.database_grounded import DatabaseGroundedFDRControl
from winnow.constants import RESIDUE_MASSES

from dataclasses import dataclass
from enum import Enum
import typer
from typing_extensions import Annotated
from typing import Union
import logging
from rich.logging import RichHandler
from pathlib import Path
import yaml
import pandas as pd


# --- Configuration ---
SEED = 42
MZ_TOLERANCE = 0.02
HIDDEN_DIM = 10
TRAIN_FRACTION = 0.1


class DataSource(Enum):
    """Source of a dataset to be used for calibration."""

    winnow = "winnow"
    instanovo = "instanovo"
    pointnovo = "pointnovo"
    mztab = "mztab"


@dataclass
class WinnowDatasetConfig:
    """Config for calibration datasets saved through `winnow`."""

    data_dir: Path


@dataclass
class InstaNovoDatasetConfig:
    """Config for calibration datasets generated by InstaNovo."""

    beam_predictions_path: Path
    spectrum_path: Path


@dataclass
class MZTabDatasetConfig:
    """Config for calibration datasets saved in MZTab format."""

    spectrum_path: Path
    predictions_path: Path


@dataclass
class PointNovoDatasetConfig:
    """Config for calibration datasets generated by PointNovo."""

    mgf_path: Path
    predictions_path: Path


class FDRMethod(Enum):
    """FDR estimation method."""

    database = "database-ground"
    winnow = "winnow"


# --- Logging Setup ---
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.addHandler(logging.StreamHandler())
logger.addHandler(RichHandler())

app = typer.Typer(
    name="winnow",
    help="""
    Confidence calibration and FDR estimation for de novo peptide sequencing.
    """,
)


def load_dataset(
    data_source: DataSource, dataset_config_path: Path
) -> CalibrationDataset:
    """Load PSM dataset into a `CalibrationDataset` object.

    Args:
        data_source (DataSource): The source of the dataset
        dataset_config_path (Path): Path to a `.yaml` file containing arguments
                                    for the load method for the data source.

    Raises:
        TypeError: If `data_source` is not one of the supported data sources

    Returns:
        CalibrationDataset: A calibration dataset
    """
    logger.info(f"Loading dataset from {data_source}.")
    with open(dataset_config_path) as dataset_config_file:
        if data_source is DataSource.winnow:
            winnow_dataset_config = WinnowDatasetConfig(
                **yaml.safe_load(dataset_config_file)
            )
            dataset = WinnowDatasetLoader().load(
                data_path=Path(winnow_dataset_config.data_dir)
            )
        elif data_source is DataSource.instanovo:
            instanovo_dataset_config = InstaNovoDatasetConfig(
                **yaml.safe_load(dataset_config_file)
            )
            dataset = InstaNovoDatasetLoader().load(
                data_path=Path(instanovo_dataset_config.spectrum_path),
                predictions_path=Path(instanovo_dataset_config.beam_predictions_path),
            )
        elif data_source is DataSource.mztab:
            mztab_dataset_config = MZTabDatasetConfig(
                **yaml.safe_load(dataset_config_file)
            )
            dataset = MZTabDatasetLoader().load(
                data_path=Path(mztab_dataset_config.spectrum_path),
                predictions_path=Path(mztab_dataset_config.predictions_path),
            )
        elif data_source is DataSource.pointnovo:
            pointnovo_dataset_config = PointNovoDatasetConfig(
                **yaml.safe_load(dataset_config_file)
            )
            dataset = PointNovoDatasetLoader().load(
                data_path=Path(pointnovo_dataset_config.mgf_path),
                predictions_path=Path(pointnovo_dataset_config.predictions_path),
            )
        else:
            raise TypeError(
                f"Data source was {data_source}. Only 'instanovo', 'mztab' and 'pointnovo' are supported."
            )
    return dataset


def filter_dataset(dataset: CalibrationDataset) -> CalibrationDataset:
    """Filter out rows whose predictions are empty or contain unsupported PSMs.

    Args:
        dataset (CalibrationDataset): The dataset to be filtered

    Returns:
        CalibrationDataset: The filtered dataset
    """
    logger.info("Filtering dataset.")
    filtered_dataset = (
        dataset.filter_entries(
            metadata_predicate=lambda row: not isinstance(row["prediction"], list)
        )
        .filter_entries(metadata_predicate=lambda row: not row["prediction"])
        .filter_entries(
            metadata_predicate=lambda row: row["precursor_charge"] > 6
        )  # Prosit-specific filtering, see https://github.com/Nesvilab/FragPipe/issues/1775
        .filter_entries(
            predictions_predicate=lambda row: len(row[0].sequence) > 30
        )  # Prosit-specific filtering
        .filter_entries(
            predictions_predicate=lambda row: len(row[1].sequence) > 30
        )  # Prosit-specific filtering
    )
    return filtered_dataset


def initialise_calibrator() -> ProbabilityCalibrator:
    """Set up the probability calibrator with features."""
    calibrator = ProbabilityCalibrator(SEED)
    calibrator.add_feature(MassErrorFeature(residue_masses=RESIDUE_MASSES))
    calibrator.add_feature(PrositFeatures(mz_tolerance=MZ_TOLERANCE))
    calibrator.add_feature(
        RetentionTimeFeature(hidden_dim=HIDDEN_DIM, train_fraction=TRAIN_FRACTION)
    )
    calibrator.add_feature(ChimericFeatures(mz_tolerance=MZ_TOLERANCE))
    calibrator.add_feature(BeamFeatures())
    return calibrator


def apply_fdr_control(
    fdr_control: Union[NonParametricFDRControl, DatabaseGroundedFDRControl],
    dataset: CalibrationDataset,
    fdr_threshold: float,
    confidence_column: str,
) -> pd.DataFrame:
    """Apply FDR control to a dataset."""
    if isinstance(fdr_control, NonParametricFDRControl):
        fdr_control.fit(dataset=dataset.metadata[confidence_column])
        dataset.metadata = fdr_control.add_psm_pep(dataset.metadata, confidence_column)
        dataset.metadata = fdr_control.add_psm_q_value(
            dataset.metadata, confidence_column
        )
    else:
        fdr_control.fit(
            dataset=dataset.metadata[confidence_column],
            residue_masses=RESIDUE_MASSES,
        )
    dataset.metadata = fdr_control.add_psm_fdr(dataset.metadata, confidence_column)
    dataset.metadata = fdr_control.add_psm_q_value(dataset.metadata, confidence_column)
    confidence_cutoff = fdr_control.get_confidence_cutoff(threshold=fdr_threshold)
    output_data = dataset.metadata
    output_data = output_data[output_data[confidence_column] >= confidence_cutoff]
    return output_data


def check_if_labelled(dataset: CalibrationDataset) -> None:
    """Check if the dataset contains a ground-truth column."""
    if "sequence" not in dataset.metadata.columns:
        raise ValueError(
            "Database-grounded FDR control can only be performed on annotated data."
        )


@app.command(name="train", help="Fit a calibration model.")
def train(
    data_source: Annotated[
        DataSource, typer.Option(help="The type of PSM dataset to be calibrated.")
    ],
    dataset_config_path: Annotated[
        Path,
        typer.Option(
            help="The path to the config with the specification of the calibration dataset."
        ),
    ],
    model_output_folder: Annotated[
        Path, typer.Option(help="The path to write the fitted model checkpoints to.")
    ],
    dataset_output_path: Annotated[
        Path, typer.Option(help="The path to write the output to.")
    ],
):
    """Fit the calibration model.

    Args:
        data_source (Annotated[ DataSource, typer.Option, optional): The type of PSM dataset to be calibrated.
        dataset_config_path (Annotated[ Path, typer.Option, optional): The path to the config with the specification of the calibration dataset.
        model_output_folder (Annotated[Path, typer.Option, optional]): The path to write the fitted model checkpoints to.
        dataset_output_path (Annotated[Path, typer.Option, optional): The path to write the output to.
    """
    # -- Load dataset
    logger.info("Loading datasets.")
    annotated_dataset = load_dataset(
        data_source=data_source,
        dataset_config_path=dataset_config_path,
    )

    annotated_dataset = filter_dataset(annotated_dataset)

    # Train
    logger.info("Training calibrator.")
    calibrator = initialise_calibrator()
    calibrator.fit(annotated_dataset)

    # -- Write model checkpoints
    logger.info(f"Saving model to {model_output_folder}")
    ProbabilityCalibrator.save(calibrator, model_output_folder)

    # -- Write output
    logger.info("Writing output.")
    annotated_dataset.to_csv(dataset_output_path)
    logger.info(f"Training dataset results saved: {dataset_output_path}")


@app.command(
    name="predict",
    help="Calibrate scores and optionally filter results to a target FDR.",
)
def predict(
    data_source: Annotated[
        DataSource, typer.Option(help="The type of PSM dataset to be calibrated.")
    ],
    dataset_config_path: Annotated[
        Path,
        typer.Option(
            help="The path to the config with the specification of the calibration dataset."
        ),
    ],
    model_folder: Annotated[
        Path, typer.Option(help="The path to calibrator checkpoints.")
    ],
    method: Annotated[
        FDRMethod, typer.Option(help="Method to use for FDR estimation.")
    ],
    fdr_threshold: Annotated[
        float,
        typer.Option(
            help="The target FDR threshold (e.g. 0.01 for 1%, 0.05 for 5% etc.)"
        ),
    ],
    confidence_column: Annotated[
        str, typer.Option(help="Name of the column with confidence scores.")
    ],
    output_path: Annotated[Path, typer.Option(help="The path to write the output to.")],
):
    """Calibrate model scores, estimate FDR and filter for a threshold.

    Args:
        data_source (Annotated[ DataSource, typer.Option, optional): The type of PSM dataset to be calibrated.
        dataset_config_path (Annotated[ Path, typer.Option, optional): The path to the config with the specification of the dataset.
        model_folder (Annotated[ Path, typer.Option, optional): The path to calibrator checkpoints.
        method (Annotated[ FDRMethod, typer.Option, optional): Method to use for FDR estimation.
        fdr_threshold (Annotated[ float, typer.Option, optional): The target FDR threshold (e.g. 0.01 for 1%, 0.05 for 5% etc.).
        confidence_column (Annotated[ str, typer.Option, optional): Name of the column with confidence scores.
        output_path (Annotated[ Path, typer.Option, optional): The path to write the output to.
    """
    # -- Load dataset
    logger.info("Loading datasets.")
    dataset = load_dataset(
        data_source=data_source,
        dataset_config_path=dataset_config_path,
    )

    dataset = filter_dataset(dataset)

    # Predict
    logger.info("Loading calibrator.")
    calibrator = ProbabilityCalibrator.load(model_folder)

    logger.info("Calibrating scores.")
    calibrator.predict(dataset)

    if method is FDRMethod.winnow:
        logger.info("Applying FDR control.")
        dataset_metadata = apply_fdr_control(
            NonParametricFDRControl(), dataset, fdr_threshold, confidence_column
        )
    elif method is FDRMethod.database:
        logger.info("Applying FDR control.")
        check_if_labelled(dataset)
        dataset_metadata = apply_fdr_control(
            DatabaseGroundedFDRControl(confidence_feature=confidence_column),
            dataset,
            fdr_threshold,
            confidence_column,
        )

    # -- Write output
    logger.info("Writing output.")
    dataset_metadata.to_csv(output_path)
    logger.info(f"Outputs saved: {output_path}")
