#!/usr/bin/env python3

#Auteur : Pierre Koclas, May 2021
import os
import sys
import csv
from math import floor,ceil,sqrt
import matplotlib as mpl
mpl.use('Agg')
#import pylab as plt
import matplotlib.pylab as plt
import numpy as np
import matplotlib.colorbar as cbar
import matplotlib.cm as cm
import datetime
import cartopy.crs as ccrs
import cartopy.feature
#from cartopy.mpl.ticker    import LongitudeFormatter,  LatitudeFormatter
import matplotlib.colors as colors
#import matplotlib.patches as mpatches
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
import sqlite3
from matplotlib.collections import PatchCollection
from statistics import median
import pikobs
import optparse
from mpl_toolkits.axes_grid1 import make_axes_locatable
def projectPpoly(PROJ,lat,lon,deltax,deltay,pc):
        X1,Y1  = PROJ.transform_point(lon - deltax,lat-deltay,pc )
        X2,Y2  = PROJ.transform_point(lon - deltax,lat+deltay,pc )
        X3,Y3  = PROJ.transform_point(lon + deltax,lat+deltay,pc )
        X4, Y4 = PROJ.transform_point(lon + deltax,lat-deltay,pc )
        Pt1=[ X1,Y1 ]
        Pt2=[ X2,Y2 ]
        Pt3=[ X3,Y3 ]
        Pt4=[ X4,Y4 ]
        Points4 = [ Pt1, Pt2,Pt3,Pt4 ]
           
        return Points4
def SURFLL(lat1,lat2,lon1,lon2):
#= (pi/180)R^2 |sin(lat1)-sin(lat2)| |lon1-lon2|
    R=6371.
    lat2=min(lat2,90.)
    surf=R*R*(np.pi/180.)*abs ( np.sin(lat2*np.pi/180.) - np.sin(lat1*np.pi/180.) ) *abs( lon2-lon1 )
   # if ( surf == 0.):
    # print (   ' surf=',lat1,lat2,lat2*np.pi/180.,lat1*np.pi/180.,np.sin(lat2*np.pi/180.) ,  np.sin(lat1*np.pi/180.) )
    return surf

def NPSURFLL(lat1, lat2, lon1, lon2):
    R = 6371.
    lat2 = np.minimum(lat2, 90.)
    surf = R**2 * (np.pi/180) * np.abs(np.sin(lat2*np.pi/180) - np.sin(lat1*np.pi/180)) * np.abs(lon2 - lon1)
  #  if np.any(surf == 0.):
    #    print('surf contiene valores cero')
    return surf
def SURFLL2(lat1, lat2, lon1, lon2):
    R = 6371.0
    lat2 = np.minimum(lat2, 90.0)
    surf = R * R * (np.pi / 180.0) * np.abs(np.sin(lat2 * np.pi / 180.0) - np.sin(lat1 * np.pi / 180.0)) * np.abs(lon2 - lon1)
    # Debugging print statements if surface is zero
    zero_surf_indices = (surf == 0.0)
    if np.any(zero_surf_indices):
        print('surf=', lat1[zero_surf_indices], lat2[zero_surf_indices], lat2[zero_surf_indices] * np.pi / 180.0,
              lat1[zero_surf_indices] * np.pi / 180.0,
              np.sin(lat2[zero_surf_indices] * np.pi / 180.0),
              np.sin(lat1[zero_surf_indices] * np.pi / 180.0))
    return surf
##           plt.close(fig)
def days_between(d1, d2):
    d1 = datetime.datetime.strptime(d1, "%Y%m%d%H")
    d2 = datetime.datetime.strptime(d2, "%Y%m%d%H")
    return abs((d2 - d1).days)

import pikobs

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import sqlite3
import os
from rich.console import Console
from rich.table import Table
from rich import box

def obscountdb_plot(
    region,
    families, 
    datestart,
    dateend, 
    pathwork, 
    namesin,
    flag_criteria
):
    if len(namesin) != 2:
        raise ValueError("The 'namesin' list must contain exactly two names to compare.")
    
    name1, name2 = namesin[0], namesin[1] # name1 is the REFERENCE (0)
    console = Console()
    
    list_df_overall = []
    list_df_stn = []

    console.print(f"[bold yellow]Extracting data from databases in: {pathwork}[/bold yellow]")

    # 1. DATA EXTRACTION
    for family in families:
        for name in namesin:
            # Check if your structure is pathwork/name_family/ or pathwork/family/. Adjust if needed.
            # Here assuming pathwork/2026020100_ai/ depending on how you pass pathwork
            db_file = os.path.join(pathwork, f"{name}_{family}", f'{name}_{region}_{datestart}_{dateend}_{flag_criteria}_{family}.db')
            
            # Fallback just in case the folder is named only by family:
            if not os.path.exists(db_file):
                db_file = os.path.join(pathwork, family, f'{name}_{region}_{datestart}_{dateend}_{flag_criteria}_{family}.db')

            if not os.path.exists(db_file):
                console.print(f"[red]  Warning: File not found {db_file}[/red]")
                continue
                
            conn = sqlite3.connect(db_file)
            
            # OVERALL TOTALS (Runs for ALL families)
            query_overall = "SELECT SUM(Nobsdata) as Nobs, SUM(Nobsprofile) as NobsProfiles FROM moyenne;"
            df_ov = pd.read_sql_query(query_overall, conn)
            df_ov['family'], df_ov['name'] = family, name
            list_df_overall.append(df_ov)
            
            # BY STATION
            if family in ['ai', 'sf', 'csr','gp','ua']:
                console.print(f"[cyan]  Skipping station breakdown for '{family}' (Too many id_stn).[/cyan]")
                
                # FUTURE OPTION: Grouping by 'LIKE' for 'ai' or 'sf'
                # Uncomment and adjust the SQL below when you want to group them by a specific pattern.
                """
                query_stn_like = '''
                    SELECT 
                        CASE 
                            WHEN id_stn LIKE 'PREFIX1%' THEN 'Group_1'
                            WHEN id_stn LIKE 'PREFIX2%' THEN 'Group_2'
                            ELSE 'Other' 
                        END as id_stn,
                        SUM(Nobsdata) as Nobsdata, 
                        SUM(Nobsprofile) as NobsProfile
                    FROM moyenne 
                    GROUP BY 1;
                '''
                df_stn = pd.read_sql_query(query_stn_like, conn)
                df_stn['family'], df_stn['name'] = family, name
                list_df_stn.append(df_stn)
                """
            else:
                # Normal processing for other families
                query_stn = "SELECT id_stn, SUM(Nobsdata) as Nobs, SUM(Nobsprofile) as NobsProfiles FROM moyenne GROUP BY id_stn;"
                df_stn = pd.read_sql_query(query_stn, conn)
                df_stn['family'], df_stn['name'] = family, name
                list_df_stn.append(df_stn)
            
            conn.close()

    # Helper function: Calculate percentage safely to avoid division by zero
    def calculate_pct(diff, ref):
        return np.where(ref == 0, np.where(diff > 0, 100.0, 0.0), (diff / ref) * 100.0)

    # Helper function: Format DataFrame to Strings for clean PNG rendering
    def format_df_for_png(df):
        df_png = df.copy()
        for col in ['Diff_Nobs', 'Diff_Profiles']:
            if col in df_png.columns:
                df_png[col] = df_png[col].apply(lambda x: f"+{int(x)}" if x > 0 else str(int(x)))
        for col in ['Pct_Nobs', 'Pct_Profiles']:
            if col in df_png.columns:
                df_png[col] = df_png[col].apply(lambda x: f"+{x:.1f}%" if x > 0 else f"{x:.1f}%")
        return df_png

    # Helper function: Save DataFrame as PNG
    def save_df_as_png(df, filename, title):
        fig_height = max(2, len(df) * 0.35 + 1.5)
        fig, ax = plt.subplots(figsize=(16, fig_height)) # Width increased
        ax.axis('off')
        ax.axis('tight')

        mpl_table = ax.table(
            cellText=df.values,
            colLabels=df.columns,
            cellLoc='center',
            loc='center',
            colColours=['#4C72B0'] * len(df.columns)
        )
        
        mpl_table.auto_set_font_size(False)
        mpl_table.set_fontsize(10)
        mpl_table.scale(1.2, 1.5)
        for k, cell in mpl_table._cells.items():
            cell.set_edgecolor('white')
            if k[0] == 0:
                cell.set_text_props(weight='bold', color='white')

        plt.title(title, fontsize=14, fontweight='bold', pad=20)
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.close()

    # Helper function: Color logic for HTML Pandas Styler
    def color_html_diff(val):
        if isinstance(val, (int, float)):
            if val > 0: return 'color: green; font-weight: bold;'
            elif val < 0: return 'color: red; font-weight: bold;'
        return ''

    # --- HTML REPORT INITIALIZATION ---
    html_content = f"""
    <html>
    <head>
        <meta charset="utf-8">
        <title>Observation Report ({region})</title>
        <style>
            body {{ font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; margin: 40px; background-color: #f9f9f9; }}
            h1, h2, h3 {{ color: #333; }}
            h1 {{ border-bottom: 2px solid #4C72B0; padding-bottom: 10px; }}
            table {{ border-collapse: collapse; width: 95%; margin-bottom: 40px; background-color: white; box-shadow: 0 2px 5px rgba(0,0,0,0.1); }}
            th, td {{ border: 1px solid #ddd; padding: 12px; text-align: center; }}
            th {{ background-color: #4C72B0; color: white; position: sticky; top: 0; }}
            tr:nth-child(even) {{ background-color: #f2f2f2; }}
            tr:hover {{ background-color: #ddd; }}
        </style>
    </head>
    <body>
        <h1>Observation Comparison: {name1} (Ref) vs {name2} | Region: {region}</h1>
    """

    cols_order = [
        'family' if 'family' in locals() else 'id_stn', 
        f'Nobs (Ref: {name1})', f'Nobs ({name2})', 'Diff_Nobs', 'Pct_Nobs', 
        f'Profiles (Ref: {name1})', f'Profiles ({name2})', 'Diff_Profiles', 'Pct_Profiles'
    ]

    # 2. PROCESS AND DISPLAY OVERALL DATA
    if list_df_overall:
        df_all = pd.concat(list_df_overall).fillna(0)
        pivot_ov = df_all.pivot(index='family', columns='name', values=['Nobs', 'NobsProfiles']).fillna(0)
        pivot_ov.columns = [f"{col[0]}_{col[1]}" for col in pivot_ov.columns]
        
        pivot_ov['Diff_Nobs'] = pivot_ov[f'Nobs_{name2}'] - pivot_ov[f'Nobs_{name1}']
        pivot_ov['Pct_Nobs'] = calculate_pct(pivot_ov['Diff_Nobs'], pivot_ov[f'Nobs_{name1}'])
        
        pivot_ov['Diff_Profiles'] = pivot_ov[f'NobsProfiles_{name2}'] - pivot_ov[f'NobsProfiles_{name1}']
        pivot_ov['Pct_Profiles'] = calculate_pct(pivot_ov['Diff_Profiles'], pivot_ov[f'NobsProfiles_{name1}'])
        
        pivot_ov = pivot_ov.rename(columns={
            f'Nobs_{name1}': f'Nobs (Ref: {name1})', f'Nobs_{name2}': f'Nobs ({name2})',
            f'NobsProfiles_{name1}': f'Profiles (Ref: {name1})', f'NobsProfiles_{name2}': f'Profiles ({name2})'
        }).reset_index()
        
        ov_cols = ['family'] + cols_order[1:]
        pivot_ov = pivot_ov[ov_cols]
        for col in [f'Nobs (Ref: {name1})', f'Nobs ({name2})', 'Diff_Nobs', f'Profiles (Ref: {name1})', f'Profiles ({name2})', 'Diff_Profiles']:
            pivot_ov[col] = pivot_ov[col].astype(int)

        table_ov = Table(title=f"OVERALL SUMMARY: {name1} (Ref) vs {name2}", box=box.ROUNDED, header_style="bold cyan")
        for col in pivot_ov.columns: table_ov.add_column(col, justify="center")

        for _, row in pivot_ov.iterrows():
            d_nobs = f"[bold green]+{row['Diff_Nobs']}[/]" if row['Diff_Nobs'] > 0 else (f"[bold red]{row['Diff_Nobs']}[/]" if row['Diff_Nobs'] < 0 else "0")
            p_nobs = f"[bold green]+{row['Pct_Nobs']:.1f}%[/]" if row['Pct_Nobs'] > 0 else (f"[bold red]{row['Pct_Nobs']:.1f}%[/]" if row['Pct_Nobs'] < 0 else "0.0%")
            
            d_prof = f"[bold green]+{row['Diff_Profiles']}[/]" if row['Diff_Profiles'] > 0 else (f"[bold red]{row['Diff_Profiles']}[/]" if row['Diff_Profiles'] < 0 else "0")
            p_prof = f"[bold green]+{row['Pct_Profiles']:.1f}%[/]" if row['Pct_Profiles'] > 0 else (f"[bold red]{row['Pct_Profiles']:.1f}%[/]" if row['Pct_Profiles'] < 0 else "0.0%")
            
            table_ov.add_row(str(row['family']), str(row[f'Nobs (Ref: {name1})']), str(row[f'Nobs ({name2})']), d_nobs, p_nobs, str(row[f'Profiles (Ref: {name1})']), str(row[f'Profiles ({name2})']), d_prof, p_prof)
        console.print(table_ov)
        
        overall_png_path = os.path.join(pathwork, f"Overall_Summary_{region}.png")
        save_df_as_png(format_df_for_png(pivot_ov), overall_png_path, f"Overall Comparison ({region})")
        console.print(f"[green]✔ Overall PNG saved to: {overall_png_path}[/green]\n")
        
        html_content += "<h2>Overall Summary</h2>"
        format_dict = {'Pct_Nobs': '{:+.1f}%', 'Pct_Profiles': '{:+.1f}%'}
        styled_ov = pivot_ov.style.applymap(color_html_diff, subset=['Diff_Nobs', 'Pct_Nobs', 'Diff_Profiles', 'Pct_Profiles']).format(format_dict).hide(axis='index')
        html_content += styled_ov.to_html()

    # 3. PROCESS AND DISPLAY STATION & FAMILY DATA
    if list_df_stn:
        df_stn_all = pd.concat(list_df_stn).fillna(0)
        pivot_stn = df_stn_all.pivot(index=['family', 'id_stn'], columns='name', values=['Nobs', 'NobsProfiles']).fillna(0)
        pivot_stn.columns = [f"{col[0]}_{col[1]}" for col in pivot_stn.columns]
        
        pivot_stn['Diff_Nobs'] = pivot_stn[f'Nobs_{name2}'] - pivot_stn[f'Nobs_{name1}']
        pivot_stn['Pct_Nobs'] = calculate_pct(pivot_stn['Diff_Nobs'], pivot_stn[f'Nobs_{name1}'])
        
        pivot_stn['Diff_Profiles'] = pivot_stn[f'NobsProfiles_{name2}'] - pivot_stn[f'NobsProfiles_{name1}']
        pivot_stn['Pct_Profiles'] = calculate_pct(pivot_stn['Diff_Profiles'], pivot_stn[f'NobsProfiles_{name1}'])
        
        pivot_stn = pivot_stn.rename(columns={
            f'Nobs_{name1}': f'Nobs (Ref: {name1})', f'Nobs_{name2}': f'Nobs ({name2})',
            f'NobsProfiles_{name1}': f'Profiles (Ref: {name1})', f'NobsProfiles_{name2}': f'Profiles ({name2})'
        }).reset_index()

        stn_cols = ['family', 'id_stn'] + cols_order[1:]
        pivot_stn = pivot_stn[stn_cols]
        for col in [f'Nobs (Ref: {name1})', f'Nobs ({name2})', 'Diff_Nobs', f'Profiles (Ref: {name1})', f'Profiles ({name2})', 'Diff_Profiles']:
            pivot_stn[col] = pivot_stn[col].astype(int)

        html_content += "<h2>Station Breakdown by Family</h2>"

        for fam in families:
            # Skip families that weren't added to list_df_stn (like 'ai' and 'sf')
            if fam not in pivot_stn['family'].values:
                continue

            df_fam = pivot_stn[pivot_stn['family'] == fam].drop(columns=['family'])
            if df_fam.empty: continue
                
            table_fam = Table(title=f"STATIONS - Family: {fam}", box=box.SIMPLE, header_style="bold magenta")
            for col in df_fam.columns: table_fam.add_column(col, justify="center")

            for _, row in df_fam.iterrows():
                d_nobs = f"[green]+{row['Diff_Nobs']}[/]" if row['Diff_Nobs'] > 0 else (f"[red]{row['Diff_Nobs']}[/]" if row['Diff_Nobs'] < 0 else "0")
                p_nobs = f"[green]+{row['Pct_Nobs']:.1f}%[/]" if row['Pct_Nobs'] > 0 else (f"[red]{row['Pct_Nobs']:.1f}%[/]" if row['Pct_Nobs'] < 0 else "0.0%")
                
                d_prof = f"[green]+{row['Diff_Profiles']}[/]" if row['Diff_Profiles'] > 0 else (f"[red]{row['Diff_Profiles']}[/]" if row['Diff_Profiles'] < 0 else "0")
                p_prof = f"[green]+{row['Pct_Profiles']:.1f}%[/]" if row['Pct_Profiles'] > 0 else (f"[red]{row['Pct_Profiles']:.1f}%[/]" if row['Pct_Profiles'] < 0 else "0.0%")
                
                table_fam.add_row(str(row['id_stn']), str(row[f'Nobs (Ref: {name1})']), str(row[f'Nobs ({name2})']), d_nobs, p_nobs, str(row[f'Profiles (Ref: {name1})']), str(row[f'Profiles ({name2})']), d_prof, p_prof)
            console.print(table_fam)
            
            stn_png_path = os.path.join(pathwork, f"Stations_{fam}_{region}.png")
            save_df_as_png(format_df_for_png(df_fam), stn_png_path, f"Station Breakdown - Family: {fam} ({region})")
            console.print(f"[green]✔ Stations PNG saved to: {stn_png_path}[/green]\n")
            
            html_content += f"<h3>Family: {fam}</h3>"
            styled_fam = df_fam.style.applymap(color_html_diff, subset=['Diff_Nobs', 'Pct_Nobs', 'Diff_Profiles', 'Pct_Profiles']).format(format_dict).hide(axis='index')
            html_content += styled_fam.to_html()

    html_content += """
    </body>
    </html>
    """
    html_filepath = os.path.join(pathwork, f"Full_Report_{region}.html")
    with open(html_filepath, "w", encoding="utf-8") as f:
        f.write(html_content)
    
    console.print(f"[bold cyan]🚀 Process finished! Web report saved to: {html_filepath}[/bold cyan]")
