"""
rescore_db.py — One-time migration to recompute composite_score for all DB rows.

Rewrites composite_score using the new geo formula:
    sqrt(Calmar × Sortino) × (1 + min(log(trades/floor), log(4))) × ln(1 + P&L% / 1000)

Run once after updating config.py with the new composite_score() signature.
Supports retroactive subperiod consistency check via --add-subperiod-check.

Usage:
    python3 tools/rescore_db.py                        # dry run — show before/after composite stats
    python3 tools/rescore_db.py --apply                # write updated composite scores to DB

    python3 tools/rescore_db.py --add-subperiod-check           # dry run — show how many would fail
    python3 tools/rescore_db.py --add-subperiod-check --apply   # write subperiod_consistent + zero composites
"""

import sqlite3
import math
import argparse
import glob
import multiprocessing
import pandas as pd

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from config import (TRAIN_END, SCORE_START, get_min_trades, MIN_WINNER_PNL_PCT,
                    COMPOSITE_LOG_CAP, composite_score, SUBPERIOD_SPLIT, mlp_data_path)

DB_FILE = "results/sweep_database.db"
MAX_WORKERS = 16

# Module-level cache: data file loaded once per worker process via pool initializer
_worker_df = None
_worker_score_start = None
_worker_split = None


def _worker_init(data_file, score_start, split):
    """Pool initializer: load and parse data CSV once per worker process."""
    global _worker_df, _worker_score_start, _worker_split
    import pandas as _pd
    import sys as _sys
    import os as _os
    _sys.path.insert(0, _os.path.join(_os.path.dirname(__file__), '..'))

    df = _pd.read_csv(data_file)
    df.columns = df.columns.str.lower()
    if 'time' in df.columns:
        df['time'] = _pd.to_datetime(df['time'], utc=True).dt.tz_localize(None)
    _worker_df = df
    _worker_score_start = score_start
    _worker_split = split


def _check_subperiod_one(args):
    """
    Worker function for parallel subperiod consistency check.
    Must be at module level for multiprocessing pickling.
    Uses _worker_df loaded by pool initializer — no large data passed per task.
    Returns (row_id, is_consistent: bool) or (row_id, None) on error.
    """
    row_id, params_dict = args
    import pandas as _pd
    import strategies.strategy_activation_scores as _strat
    try:
        df_res = _strat.generate_signals(_worker_df.copy(), **params_dict)

        split_ts = _pd.to_datetime(_worker_split)
        m1 = _strat.calculate_metrics(
            df_res[df_res['time'] <= split_ts].copy(), score_start=_worker_score_start)
        m2 = _strat.calculate_metrics(
            df_res[df_res['time'] > split_ts].copy(), score_start=_worker_split)

        is_consistent = bool(
            m1.get('Calmar Ratio', -10.0) > 0 and
            m2.get('Calmar Ratio', -10.0) > 0
        )
        return (row_id, is_consistent)
    except Exception:
        return (row_id, None)


def compute_is_years(asset, timeframe):
    """Derive is_years for an asset/TF from its data CSV (calendar days in scoring window)."""
    data_path = mlp_data_path(asset, timeframe)
    candidates = [data_path] if os.path.exists(data_path) else []
    if not candidates:
        return 7.75  # fallback
    df = pd.read_csv(candidates[0])
    df.columns = df.columns.str.lower()
    if 'time' not in df.columns:
        return 7.75
    df['time'] = pd.to_datetime(df['time'], utc=True).dt.tz_localize(None)
    df_score = df[(df['time'] >= pd.to_datetime(SCORE_START)) & (df['time'] <= pd.to_datetime(TRAIN_END))]
    if len(df_score) < 2:
        return 7.75
    days = (df_score['time'].iloc[-1] - df_score['time'].iloc[0]).days
    return max(0.5, days / 365.25)


def run_subperiod_check(apply: bool):
    """Retroactively check IS subperiod consistency for all rows with composite_score > 0."""
    conn = sqlite3.connect(DB_FILE)

    # Ensure subperiod_consistent column exists (migration)
    existing = {r[1] for r in conn.execute("PRAGMA table_info(sweep_results)").fetchall()}
    if 'subperiod_consistent' not in existing:
        conn.execute("ALTER TABLE sweep_results ADD COLUMN subperiod_consistent INTEGER")
        conn.commit()
        print("Migrated: added subperiod_consistent column to sweep_results")

    # Get param column names
    col_info = conn.execute("PRAGMA table_info(sweep_results)").fetchall()
    all_cols = [r[1] for r in col_info]
    param_cols = [c for c in all_cols if c.startswith('i_')]

    # Query unchecked rows with composite_score > 0
    cursor = conn.execute(
        "SELECT * FROM sweep_results WHERE composite_score > 0 AND subperiod_consistent IS NULL"
    )
    col_names = [d[0] for d in cursor.description]
    rows = [dict(zip(col_names, r)) for r in cursor.fetchall()]
    conn.close()

    if not rows:
        print("No unchecked rows with composite_score > 0 found. Nothing to do.")
        return

    print(f"Found {len(rows):,} rows to check (composite > 0, subperiod_consistent IS NULL)")

    # Group rows by (asset, timeframe)
    from collections import defaultdict
    by_combo = defaultdict(list)
    for row in rows:
        key = (row['asset'], row['timeframe'])
        by_combo[key].append(row)

    all_results = {}  # row_id -> is_consistent (bool or None)

    for (asset, tf), combo_rows in by_combo.items():
        data_file = mlp_data_path(asset, tf)
        if not os.path.exists(data_file):
            print(f"  WARNING: no data file for {asset} {tf} — skipping {len(combo_rows)} rows")
            continue
        print(f"\n  {asset} {tf}: {len(combo_rows):,} rows — loading {data_file}")

        # Build slim args: only (row_id, params_dict) — data loaded once per worker via initializer
        worker_args = []
        for row in combo_rows:
            params_dict = {}
            for col in param_cols:
                val = row.get(col)
                if val is not None:
                    if col.startswith('i_use_'):
                        params_dict[col] = bool(int(val))
                    else:
                        params_dict[col] = float(val)
            worker_args.append((row['id'], params_dict))

        done = 0
        consistent = 0
        failed = 0
        errors = 0

        with multiprocessing.Pool(
            processes=MAX_WORKERS,
            initializer=_worker_init,
            initargs=(data_file, SCORE_START, SUBPERIOD_SPLIT)
        ) as pool:
            for row_id, is_consistent in pool.imap_unordered(_check_subperiod_one, worker_args, chunksize=50):
                all_results[row_id] = is_consistent
                done += 1
                if is_consistent is True:
                    consistent += 1
                elif is_consistent is False:
                    failed += 1
                else:
                    errors += 1
                if done % 5000 == 0 or done == len(combo_rows):
                    pct = 100 * done / len(combo_rows)
                    print(f"    {done:,}/{len(combo_rows):,} ({pct:.1f}%)  "
                          f"consistent={consistent:,}  failed={failed:,}  errors={errors:,}")

    consistent_ids = [rid for rid, v in all_results.items() if v is True]
    failed_ids     = [rid for rid, v in all_results.items() if v is False]
    error_ids      = [rid for rid, v in all_results.items() if v is None]

    total = len(all_results)
    print(f"\nResults ({total:,} rows checked):")
    print(f"  Consistent (both halves +ve Calmar): {len(consistent_ids):,}  ({100*len(consistent_ids)/max(total,1):.1f}%)")
    print(f"  Failed (one or both halves -ve):     {len(failed_ids):,}  ({100*len(failed_ids)/max(total,1):.1f}%)")
    print(f"  Errors (backtest raised exception):  {len(error_ids):,}")
    print(f"  Composite would be zeroed:           {len(failed_ids):,}")

    if not apply:
        print("\nDry run — no changes written. Re-run with --apply to update the DB.")
        return

    conn = sqlite3.connect(DB_FILE)
    cur = conn.cursor()

    if consistent_ids:
        cur.executemany(
            "UPDATE sweep_results SET subperiod_consistent=1 WHERE id=?",
            [(rid,) for rid in consistent_ids]
        )
    if failed_ids:
        cur.executemany(
            "UPDATE sweep_results SET subperiod_consistent=0, composite_score=0.0 WHERE id=?",
            [(rid,) for rid in failed_ids]
        )

    conn.commit()
    conn.close()
    print(f"\nDone. Wrote subperiod_consistent for {len(consistent_ids) + len(failed_ids):,} rows. "
          f"Zeroed composite for {len(failed_ids):,} failed rows.")


def main():
    parser = argparse.ArgumentParser(description="Rescore sweep DB with new composite formula")
    parser.add_argument("--apply", action="store_true",
                        help="Write changes to DB (default: dry run)")
    parser.add_argument("--add-subperiod-check", action="store_true",
                        help="Retroactively check IS subperiod consistency for composite>0 rows "
                             "(rows not yet checked only). Zeros composite for failures.")
    args = parser.parse_args()

    if not os.path.exists(DB_FILE):
        print(f"DB not found: {DB_FILE}")
        return

    if args.add_subperiod_check:
        run_subperiod_check(apply=args.apply)
        return

    # --- Composite re-score path ---
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute("SELECT id, asset, timeframe, calmar_ratio, sortino_ratio, total_trades, total_pnl_pct, composite_score FROM sweep_results")
    rows = cursor.fetchall()
    df = pd.DataFrame(rows, columns=['rowid', 'asset', 'timeframe', 'calmar_ratio', 'sortino_ratio', 'total_trades', 'total_pnl_pct', 'composite_score'])
    print(f"Loaded {len(df):,} rows from DB")

    # Compute is_years once per asset/TF combo
    combo_is_years = {}
    for (asset, tf) in df[['asset', 'timeframe']].drop_duplicates().itertuples(index=False):
        iy = compute_is_years(asset, tf)
        combo_is_years[(asset, tf)] = iy
        floor = get_min_trades(iy)
        print(f"  {asset} {tf}: is_years={iy:.2f}  trade_floor={floor}")

    # Recompute composite for every row
    new_scores = []
    for _, row in df.iterrows():
        iy = combo_is_years.get((row['asset'], row['timeframe']), 7.75)
        new_scores.append(composite_score(
            float(row['calmar_ratio']),
            float(row['sortino_ratio']),
            int(row['total_trades']),
            iy,
            float(row['total_pnl_pct'])
        ))
    df['new_composite'] = new_scores

    # Summary stats
    old_pos = (df['composite_score'] > 0).sum()
    new_pos = (df['new_composite'] > 0).sum()
    print(f"\nBefore: mean={df['composite_score'].mean():.4f}  max={df['composite_score'].max():.4f}  rows>0={old_pos:,}")
    print(f"After:  mean={df['new_composite'].mean():.4f}  max={df['new_composite'].max():.4f}  rows>0={new_pos:,}")
    print(f"Rows disqualified by P&L floor: {old_pos - new_pos:,}")

    # Per-combo top result before/after
    print(f"\n{'Combo':<35}  {'Old top Composite':>18}  {'New top Composite':>18}  {'Old P&L%':>12}  {'New P&L%':>12}")
    print("-" * 100)
    for (asset, tf), grp in df.groupby(['asset', 'timeframe']):
        old_best = grp.loc[grp['composite_score'].idxmax()]
        new_best = grp.loc[grp['new_composite'].idxmax()]
        combo = f"{asset} {tf}"
        print(f"{combo:<35}  {old_best['composite_score']:>18.4f}  {new_best['new_composite']:>18.4f}  "
              f"{old_best['total_pnl_pct']:>12,.0f}  {new_best['total_pnl_pct']:>12,.0f}")

    if not args.apply:
        print("\nDry run — no changes written. Re-run with --apply to update the DB.")
        conn.close()
        return

    # Write back
    cursor = conn.cursor()
    for _, row in df.iterrows():
        cursor.execute("UPDATE sweep_results SET composite_score=? WHERE id=?",
                       (row['new_composite'], int(row['rowid'])))
    conn.commit()
    conn.close()
    print(f"\nDone. Updated {len(df):,} rows.")


if __name__ == "__main__":
    main()