#!/usr/bin/python3
"""
Description
------------
This module calculates and plots the error histogram of observations over a given period.
"""
import sqlite3
import re
import os
import glob
from datetime import datetime, timedelta
from pathlib import Path
import argparse
import sys
import time
import numpy as np
import dask
from dask.distributed import Client

# Custom local imports
import pikobs

# --- Constants ---
ALL_STATIONS = 'all'
JOIN_CHANNELS = 'join'

# --- Database and Aggregation Functions ---

# <--- CAMBIO: La función ahora recibe 'selected_flags'
def create_and_insert_flag_observations(input_file: str, output_file: Path, region: str, selected_flags: str):
    """
    Reads an input SQLite file, filters data by region and flags, aggregates statistics,
    and inserts them into a central output database.
    """
    basename = Path(input_file).name
    try:
        date_part = basename.split('_')[0]
    except IndexError:
        print(f"Warning: Could not parse date from filename '{basename}'. Skipping.")
        return

    with sqlite3.connect(input_file) as conn_input:
        cursor_input = conn_input.cursor()

        LON1, LON2, LAT1, LAT2 = pikobs.regions(region)
        lat_lon_crit = pikobs.generate_latlon_criteria(LAT1, LAT2, LON1, LON2)
        flag_criteria = pikobs.flag_criteria(selected_flags).strip().removeprefix("and").strip()# <--- AHORA USA EL ARGUMENTO

        query = f"""
            SELECT ?, d.flag, COUNT(*), h.id_stn, d.varno, d.vcoord
            FROM data AS d
            JOIN header AS h ON d.id_obs = h.id_obs
            WHERE {flag_criteria}  {lat_lon_crit}
            GROUP BY h.date, d.varno, d.flag, h.id_stn, d.vcoord;
        """
        print (query)
        cursor_input.execute(query, (date_part,))
        results = cursor_input.fetchall()

    if not results:
        return

    with sqlite3.connect(str(output_file), isolation_level=None, timeout=300) as conn_output:
        conn_output.execute('PRAGMA synchronous = OFF')
        conn_output.execute('PRAGMA journal_mode = WAL')
        cursor_output = conn_output.cursor()
        
        cursor_output.execute('''
            CREATE TABLE IF NOT EXISTS flag_observations (
                date TEXT, flag INTEGER, num_data_entries INTEGER,
                id_stn TEXT, varno INTEGER, vcoord INTEGER
            )''')
        
        cursor_output.executemany('INSERT INTO flag_observations VALUES (?, ?, ?, ?, ?, ?)', results)

# --- Plotting Task Generation ---
def create_plot_tasks(pathwork: Path, region: str, id_stn: list, channel: list) -> list:
    """Generates a list of dictionaries, each representing a plot to be created."""
    db_file = pathwork / f"output_file{region}.db"
    if not db_file.exists(): return []

    plot_tasks = []
    with sqlite3.connect(db_file) as conn:
        cursor = conn.cursor()
        cursor.execute("SELECT DISTINCT varno FROM flag_observations")
        all_varnos = [v[0] for v in cursor.fetchall()]

        for varno in all_varnos:
            if id_stn[0] == ALL_STATIONS:
                cursor.execute("SELECT DISTINCT id_stn FROM flag_observations WHERE varno = ?", (varno,))
                station_list = [s[0] for s in cursor.fetchall()]
            else:
                station_list = id_stn

            if channel[0] == JOIN_CHANNELS:
                for station_id in station_list:
                    plot_tasks.append({'id_stn': station_id, 'vcoord': JOIN_CHANNELS, 'varno': varno})
            else:
                query = "SELECT DISTINCT vcoord FROM flag_observations WHERE varno = ? AND id_stn = ?"
                for station_id in station_list:
                    cursor.execute(query, (varno, station_id))
                    for vcoord_tuple in cursor.fetchall():
                        plot_tasks.append({'id_stn': station_id, 'vcoord': vcoord_tuple[0], 'varno': varno})
    return plot_tasks

# --- Utility and Main Functions ---
def filter_files_by_date(path: Path, family: str, datestart: str, dateend: str) -> list[str]:
    """Filters files in a directory that match the family and fall within a date range."""
    date_format = "%Y%m%d%H"
    start_dt = datetime.strptime(datestart, date_format)
    end_dt = datetime.strptime(dateend, date_format)
    
    filtered_files = [str(fp) for fp in path.glob(f'*_{family}') if start_dt <= datetime.strptime(fp.name.split('_')[0], date_format) <= end_dt]
    return filtered_files

def _run_jobs(jobs: list, n_cpu: int):
    """Helper function to run a list of Dask delayed jobs either in series or parallel."""
    if not jobs: return
    
    if n_cpu > 1:
        print(f"Executing {len(jobs)} jobs in parallel on {n_cpu} CPUs...")
        with Client(processes=True, threads_per_worker=1, n_workers=n_cpu, silence_logs=40) as client:
            dask.compute(*jobs)
    else:
        print(f"Executing {len(jobs)} jobs in series...")
        dask.compute(*jobs, scheduler='single-threaded')

# <--- CAMBIO: La función ahora recibe 'selected_flags'
def make_histograms(path_experience_files: str,
                    pathwork: str,
                    datestart: str,
                    dateend: str,
                    regions: list[str],
                    familys: list[str],
                    id_stn: list[str],
                    channel: list[str],
                    plot_title: str,
                    n_cpu: int,
                    selected_flags: str):
    """Main orchestrator function."""
    path_exp = Path(path_experience_files)
    path_work = Path(pathwork)
    
    for family in familys:
        pikobs.delete_create_folder(path_work, family)
        for region in regions:
            print(f"\n--- Processing Family: {family}, Region: {region} ---")
            
            output_db = path_work / f"output_file{region}.db"
            if output_db.exists(): output_db.unlink()

            file_list = filter_files_by_date(path_exp, family, datestart, dateend)
            if not file_list:
                print(f"Warning: No source files found for family '{family}' in date range.")
                continue

            t0 = time.time()
            # <--- CAMBIO: Pasamos 'selected_flags' a la función de Dask
            aggregation_jobs = [
                dask.delayed(create_and_insert_flag_observations)(f, output_db, region, selected_flags)
                for f in file_list
            ]
            _run_jobs(aggregation_jobs, n_cpu)
            print(f"Data aggregation finished in {time.time() - t0:.2f} seconds.")

            plot_tasks = create_plot_tasks(path_work, region, id_stn, channel)
            if not plot_tasks:
                print("Warning: No plots to generate based on criteria.")
                continue

            t0 = time.time()
            plot_jobs = [
                dask.delayed(pikobs.flags_plot)(
                    str(path_work), datestart, dateend, family, plot_title,
                    task['vcoord'], task['id_stn'], task['varno'], region
                ) for task in plot_tasks
            ]
            _run_jobs(plot_jobs, n_cpu)
            print(f"Plot generation finished in {time.time() - t0:.2f} seconds.")

def main():
    """Parses command-line arguments and runs the main histogram generation."""
    parser = argparse.ArgumentParser(description="Generate error histograms for PikObs observations.")
    parser.add_argument('--path_experience_files', required=True, type=str, help="Directory of input sqlite files")
    parser.add_argument('--pathwork', required=True, type=str, help="Working directory")
    parser.add_argument('--datestart', required=True, type=str, help="Start date (YYYYMMDDHH)")
    parser.add_argument('--dateend', required=True, type=str, help="End date (YYYYMMDDHH)")
    parser.add_argument('--flags_criteria', required=True, type=str, help="Flag criteria (e.g., 'assimilee')")
    parser.add_argument('--region', nargs="+", required=True, help="List of regions")
    parser.add_argument('--family', nargs="+", required=True, help="List of families")
    parser.add_argument('--id_stn', nargs="+", default=[ALL_STATIONS], help="List of station IDs or 'all'")
    parser.add_argument('--channel', nargs="+", default=[ALL_STATIONS], help="List of channels, 'all', or 'join'")
    parser.add_argument('--plot_title', default='PikObs Histogram', type=str, help="Base title for plots")
    parser.add_argument('--n_cpus', default=1, type=int, help="Number of CPUs for parallel processing")
    args = parser.parse_args()

    print("--- Starting Histogram Generation ---")
    for arg, value in vars(args).items():
        print(f"  --{arg}: {value}")
    print("-------------------------------------")

    # <--- CAMBIO: Pasamos el nuevo argumento a la función principal
    make_histograms(
        path_experience_files=args.path_experience_files,
        pathwork=args.pathwork,
        datestart=args.datestart,
        dateend=args.dateend,
        regions=args.region,
        familys=args.family,
        id_stn=args.id_stn,
        channel=args.channel,
        plot_title=args.plot_title,
        n_cpu=args.n_cpus,
        selected_flags=args.flags_criteria
    )
    return 0

if __name__ == '__main__':
    sys.exit(main())
