"""Classifier-adjusted density estimation."""

import numpy as np
import pandas as pd
from pandahandler import Schema, categorize_non_numerics
from sklearn import metrics

from empdens import models
from empdens.base import AbstractDensity
from empdens.classifiers.base import AbstractLearner
from empdens.classifiers.lightgbm import Lgbm
from empdens.data import CadeData


def auc(df):
    """Compute the area under the ROC curve."""
    fpr, tpr, _ = metrics.roc_curve(df.truth.values, df.pred.values, pos_label=1)
    return metrics.auc(fpr, tpr)


class Cade(AbstractDensity):
    """Classifier-adjusted density estimation.

    Based on https://pdfs.semanticscholar.org/e4e6/033069a8569ba16f64da3061538bcb90bec6.pdf
    """

    # A soft target for the number of instances to simulate when `sim_size` is "auto"
    simulation_size_attractor = 10000

    def __init__(
        self,
        initial_density: AbstractDensity | None = None,
        classifier: AbstractLearner | None = None,
        sim_size: str | int = "auto",
        verbose: bool = False,
    ):
        """Initialize the classifier-adjusted density estimation model.

        Args:
            initial_density: A naive density model to use as the basis for CADE estimation. CADE uses this as a
                reference such that the final estimated density is a product of the initial density and the classifier
                adjustment. If None, defaults to a JointDensity model.
            classifier: A classifier to use for the adjustment. If None, defaults to a LightGBM classifier.
            sim_size: The number of synthetic samples to simulate. If "auto", the simulation size is set as the
                geometric mean between the data size and `simulation_size_attractor`. If a positive number less than
                100, the simulation size is `round(sim_size)*df.shape[0]`. If `sim_size` is greater than or equal to
                100, the simulation size is `round(sim_size)`.
            verbose: Whether to print diagnostic information.
        """
        super().__init__()
        self.initial_density = initial_density or models.JointDensity(verbose=verbose)
        self.classifier = classifier or Lgbm()
        self.sim_size = sim_size
        self.verbose = verbose

    def compute_simulation_size(self, df):
        """Determine the number of synthetic data samples to simulate.

        If self.sim_size is 'auto', sets the simulation size as the geometric mean
        between the data size and self.simulation_size_attractor

        If self.sim_size is a positive number less than 100, simulation size is
        round(self.sim_size)*df.shape[0]

        Finally, if self.sim_size >= 100, simulation size is round(self.sim_size)
        """
        n_real = df.shape[0]
        if isinstance(self.sim_size, str):
            assert self.sim_size == "auto"
            sim_n = np.sqrt(n_real * self.simulation_size_attractor)
        elif self.sim_size < 100:
            assert self.sim_size > 0
            sim_n = round(self.sim_size * n_real)
            if sim_n < 10:
                raise Exception("Simulation size is very small. Consider using a larger value of sim_size")
        else:
            sim_n = round(self.sim_size)
        self.sim_rate = sim_n / df.shape[0]
        return int(sim_n)

    def _diagnostics(self, x, truth):
        val_df = pd.DataFrame({"pred": self.classifier.predict(x), "truth": truth})
        self.diagnostics = {
            "val_df": val_df,
            "auc": auc(val_df),
        }

    def train(self, df: pd.DataFrame, diagnostics: bool = False):
        """Model the density of the data.

        :param df: (pandas DataFrame)
        """
        df = categorize_non_numerics(df)
        self.schema = Schema.from_df(df)
        self.vp(f"Training a generative density model on {len(df)} samples")
        self.initial_density.train(df)
        sim_n = self.compute_simulation_size(df)
        self.vp(f"Simulating {sim_n} fake samples from the model and join it with the real data")
        sim_df = self.initial_density.rvs(sim_n)
        sim_df = self.schema(sim_df)
        xdf = pd.concat([df, sim_df], axis=0)
        assert isinstance(xdf, pd.DataFrame), "CadeData.X requires a pandas DataFrame"
        partially_synthetic_data = CadeData(
            X=xdf,
            y=np.concatenate([np.ones(df.shape[0]), np.zeros(sim_n)]),
        )
        self.vp("Train the classifier to distinguish real from fake")
        self.classifier.train(partially_synthetic_data)
        if diagnostics:
            self._diagnostics(partially_synthetic_data.X, partially_synthetic_data.y)
        if self.verbose:
            if not hasattr(self, "diagnostics"):
                self._diagnostics(partially_synthetic_data.X, partially_synthetic_data.y)
            AUROC = str(round(self.diagnostics["auc"], 3))
            print("In-sample, the classifier had AUROC = " + AUROC)

    def density(self, X):
        """Predict the density at new points.

        Apply equation 2.1 in https://pdfs.semanticscholar.org/e4e6/033069a8569ba16f64da3061538bcb90bec6.pdf

        Args:
            X: Data frame matching the schema of the training data.
        """
        X = self.schema(X)
        # Initial density estimate
        synthetic_dens = self.initial_density.density(X)
        # Classifier adjustment factor
        p_real = self.classifier.predict(X)
        odds_real = p_real / (1 - p_real)
        classifier_adjustment = self.sim_rate * odds_real
        return synthetic_dens * classifier_adjustment
