#!/usr/bin/env python3

# Author : Pierre Koclas, May 2021
# Modified & Refactored for Pikobs + RMS significance

import os
import sys
import datetime
import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
import sqlalchemy
import pikobs

mpl.use('Agg')


def days_between(d1, d2):
    d1_obj = datetime.datetime.strptime(d1, "%Y%m%d%H")
    d2_obj = datetime.datetime.strptime(d2, "%Y%m%d%H")
    return abs((d2_obj - d1_obj).days)


def profile_plot(
    mode, region, family, id_stn, datestart, dateend, Points,
    boxsizex, boxsizey, proj, pathwork, flag_criteria, fonction,
    vcoord, filesin, namesin, varno, intervales
):

    selected_flags = pikobs.flag_criteria(flag_criteria)

    element = varno
    sat = id_stn

    def load_data(dbfile):
        engine = sqlalchemy.create_engine(f"sqlite:///{dbfile}")

        query = f"""
        SELECT 
            varno,
            round(vcoord) AS vcoord,
            SUM(sumy) * 1.0 / SUM(N) AS bias,
            SQRT(
                SUM(sumy2)*1.0/SUM(N) - POWER(SUM(sumy), 2)*1.0/(SUM(N)*SUM(N))
            ) AS scartype,
            MIN(sumy) AS min_sumy,
            MAX(sumy) AS max_sumy,
            SUM(N) AS total_n
        FROM moyenne
        WHERE 
            varno = {element}
            AND id_stn IN ('{sat}')
            AND sumy IS NOT NULL
        GROUP BY varno, round(vcoord)
        HAVING COUNT(*) > 1
        ORDER BY round(vcoord)
        """
        return pd.read_sql_query(query, engine)

    # --- Load Data ---
    df0 = load_data(filesin[0])
    has_exp = len(filesin) > 1 and os.path.exists(filesin[1])
    df1 = load_data(filesin[1]) if has_exp else pd.DataFrame()

    if df0.empty and df1.empty:
        print(f"⚠️ No data found for {sat} | varno {element}.")
        return

    fig, ax = plt.subplots(figsize=(14, 10))

    def plot_curve(df, label, color, alpha_val, marker_bias, marker_std):
        if df.empty:
            return

        if len(df) > 1:
            ax.fill_betweenx(df['vcoord'], df['min_sumy'], df['max_sumy'],
                             color=color, alpha=alpha_val)

        ax.plot(df['bias'], df['vcoord'], label=f'{label} (Bias)',
                color=color, linewidth=2, marker=marker_bias)

        ax.plot(df['scartype'], df['vcoord'], label=f'{label} (Std Dev)',
                color=color, linestyle='--', marker=marker_std)

    plot_curve(df0, namesin[0], 'blue', 0.07, 'o', 's')
    if has_exp:
        plot_curve(df1, namesin[1], 'red', 0.07, 'o', 's')

    ax.set_xlabel('Value')
    ax.axvline(x=0, color='black')

    if family in ('sw', 'ua', 'ai'):
        ax.set_ylabel('Pressure [hPa]')
    else:
        ax.set_ylabel('Channel')

    ax.invert_yaxis()
    ax.grid(True)

    ax.set_title(f'Profile {datestart} to {dateend}')

    # --- TEXT AREA ---
    ax2 = ax.twinx()
    ax2.set_ylim(ax.get_ylim())
    ax2.set_yticks([])

    xmin, xmax = ax.get_xlim()
    ax.set_xlim(xmin, xmax + (xmax - xmin) * 0.4)

    transform = ax.get_yaxis_transform()

    col1_x = 1.02
    col2_x = 1.15
    col3_x = 1.28

    ax.text(col1_x, 1.01, namesin[0], transform=ax.transAxes, color='blue')
    if has_exp:
        ax.text(col2_x, 1.01, namesin[1], transform=ax.transAxes, color='red')
        ax.text(col3_x, 1.01, 'Diff RMS', transform=ax.transAxes)

    # Control counts
    for _, row in df0.iterrows():
        ax.text(col1_x, row['vcoord'], f"{int(row['total_n'])}",
                transform=transform, color='blue')

    # --- SIGNIFICANCE USING RMS ---
    if has_exp and not df1.empty:
        df_comp = pd.merge(df0, df1, on='vcoord', suffixes=('_c', '_e'))

        added_better = False
        added_worse = False

        Z_THRESHOLD = 2.0   # puedes subir a 2.58 (99%)

        for _, row in df_comp.iterrows():
            vc = row['vcoord']

            sc = row['scartype_c']
            se = row['scartype_e']
            nc = row['total_n_c']
            ne = row['total_n_e']

            bc = row['bias_c']
            be = row['bias_e']

            # RMS
            rmsc = np.sqrt(bc**2 + sc**2)
            rmse = np.sqrt(be**2 + se**2)

            ax.text(col2_x, vc, f"{int(ne)}",
                    transform=transform, color='red')

            if nc > 1 and ne > 1 and rmsc > 0 and rmse > 0:

                variance_error = (rmsc**2 / (2 * nc)) + (rmse**2 / (2 * ne))

                if variance_error > 0:
                    z = abs(rmsc - rmse) / np.sqrt(variance_error)

                    diff_pct = ((rmsc - rmse) / rmsc) * 100.0

                    if z >= Z_THRESHOLD:

                        if rmse < rmsc:
                            lbl = "Better RMS" if not added_better else ""
                            ax.plot(col3_x, vc, marker='v', color='red',
                                    transform=transform, label=lbl)
                            ax.text(col3_x + 0.02, vc, f"{diff_pct:+.1f}%",
                                    transform=transform, color='red')
                            added_better = True
                        else:
                            lbl = "Worse RMS" if not added_worse else ""
                            ax.plot(col3_x, vc, marker='^', color='blue',
                                    transform=transform, label=lbl)
                            ax.text(col3_x + 0.02, vc, f"{diff_pct:+.1f}%",
                                    transform=transform, color='blue')
                            added_worse = True
                    else:
                        ax.text(col3_x, vc, "-", transform=transform, color='gray')

    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.05), ncol=4)

    plt.tight_layout()

    output_dir = os.path.join(pathwork, family)
    os.makedirs(output_dir, exist_ok=True)

    output_file = os.path.join(
        output_dir,
        f"{fonction}_id_stn_{id_stn}_{region}_varno{varno}.png"
    )

    plt.savefig(output_file, dpi=150, bbox_inches='tight')
    plt.close()

    print(f"Saved: {output_file}")
