"""
Ce module contient les fonctions de calcul du PSI (Population Stability Index) pour les variables d'entrée utilisées dans le modèle de prédiction de rendement agricole.
Il inclut :
- get_input_drift_summary() : Calcule un résumé des statistiques descriptives pour les variables d'entrée à partir des données de production.
- get_numeric_psi() : Calcule le PSI pour les variables d'entrée numériques en comparant la distribution de référence d'entraînement avec la distribution des données reçues en production.
- get_distribution_comparison(feature_name) : Compare la distribution d'une variable d'entrée spécifique entre
la référence d'entraînement et les données reçues en production, et retourne les valeurs pour les deux distributions.
"""

import math
from src.api.monitoring.db import get_connection



def get_input_drift_summary():
    """Calcule un résumé des statistiques descriptives pour les variables d'entrée
    à partir des données de production.
    Returns:
        dict: Un dictionnaire contenant les statistiques descriptives pour chaque variable d'entrée.
    """
    query = """
    SELECT
        input_json
    FROM api_requests
    WHERE input_json IS NOT NULL;
    """

    with get_connection() as conn:
        with conn.cursor() as cursor:
            cursor.execute(query)
            rows = cursor.fetchall()

    inputs = [row[0] for row in rows]

    numeric_fields = ["temp", "rainfall_mm", "pesticides", "trend"]

    summary = {}

    for field in numeric_fields:
        values = [
            float(input_data[field])
            for input_data in inputs
            if field in input_data and input_data[field] is not None
        ]

        if values:
            summary[field] = {
                "count": len(values),
                "mean": round(sum(values) / len(values), 2),
                "min": round(min(values), 2),
                "max": round(max(values), 2),
            }
        else:
            summary[field] = {
                "count": 0,
                "mean": None,
                "min": None,
                "max": None,
            }

    return summary

import math

def _calculate_psi(reference_values, production_values, bins=10):
    """
    Calcule le PSI entre une distribution de référence
    et une distribution de production.
    """

    if not reference_values or not production_values:
        return None

    min_value = min(reference_values)
    max_value = max(reference_values)

    if min_value == max_value:
        return 0

    bin_width = (max_value - min_value) / bins

    psi = 0

    for i in range(bins):
        lower = min_value + i * bin_width
        upper = min_value + (i + 1) * bin_width

        if i == bins - 1:
            ref_count = sum(lower <= value <= upper for value in reference_values)
            prod_count = sum(lower <= value <= upper for value in production_values)
        else:
            ref_count = sum(lower <= value < upper for value in reference_values)
            prod_count = sum(lower <= value < upper for value in production_values)

        ref_pct = ref_count / len(reference_values)
        prod_pct = prod_count / len(production_values)

        ref_pct = max(ref_pct, 0.0001)
        prod_pct = max(prod_pct, 0.0001)

        psi += (prod_pct - ref_pct) * math.log(prod_pct / ref_pct)

    return round(psi, 4)


def _get_psi_status(psi):
    if psi is None:
        return "Données insuffisantes"

    if psi < 0.1:
        return "Stable"

    if psi < 0.2:
        return "Légère dérive"

    return "Drift important"


def get_numeric_psi():
    features = ["temp", "rainfall_mm", "pesticides", "trend"]

    results = []

    with get_connection() as conn:
        with conn.cursor() as cursor:

            for feature in features:
                cursor.execute(
                    """
                    SELECT feature_value
                    FROM reference_distributions
                    WHERE feature_name = %s
                      AND source = 'train';
                    """,
                    (feature,)
                )

                reference_values = [
                    float(row[0])
                    for row in cursor.fetchall()
                    if row[0] is not None
                ]

                cursor.execute(
                    """
                    SELECT input_json
                    FROM api_requests
                    WHERE input_json IS NOT NULL;
                    """
                )

                production_inputs = cursor.fetchall()

                production_values = []

                for row in production_inputs:
                    input_data = row[0]

                    if feature in input_data and input_data[feature] is not None:
                        production_values.append(float(input_data[feature]))

                psi = _calculate_psi(reference_values, production_values)

                results.append(
                    {
                        "feature": feature,
                        "psi": psi,
                        "status": _get_psi_status(psi),
                        "reference_count": len(reference_values),
                        "production_count": len(production_values),
                    }
                )

    return results

def get_distribution_comparison(feature_name):
    """
    Compare la distribution d'une variable entre :
    - la référence d'entraînement
    - les données reçues en production
    """

    allowed_features = [
        "temp",
        "rainfall_mm",
        "pesticides",
        "trend",
    ]

    if feature_name not in allowed_features:
        raise ValueError(
            f"Variable non autorisée : {feature_name}"
        )

    with get_connection() as conn:
        with conn.cursor() as cursor:

            cursor.execute(
                """
                SELECT feature_value
                FROM reference_distributions
                WHERE feature_name = %s
                  AND source = 'train';
                """,
                (feature_name,)
            )

            reference_values = [
                float(row[0])
                for row in cursor.fetchall()
                if row[0] is not None
            ]

            cursor.execute(
                """
                SELECT input_json
                FROM api_requests
                WHERE input_json IS NOT NULL;
                """
            )

            production_inputs = cursor.fetchall()

    production_values = []

    for row in production_inputs:
        input_data = row[0]

        if (
            feature_name in input_data
            and input_data[feature_name] is not None
        ):
            production_values.append(
                float(input_data[feature_name])
            )

    return {
        "feature": feature_name,
        "reference": reference_values,
        "production": production_values,
    }