#!/usr/bin/env python3
"""
build_dataset.py — Build a feature+label dataset for NN training
from existing strategy signals and indicator CSV files.

Output: data/dataset.csv with columns:
  - instrument, timestamp
  - {strategy}_signal (5×: 1=CALL, -1=PUT, 0=HOLD)
  - {strategy}_confidence (5×: 0-100)
  - RSI14, ATR14, MACD_hist, body_pct, hour, day_of_week
  - target_1h .. target_10h (ternary: 1=CALL won, -1=PUT won, 0=no bet/lost)
  - best_expiry (0-10, optimal expiry for the consensus signal)

No look-ahead bias: all targets use only future close prices.
"""

import os
import sys
import numpy as np
import pandas as pd

# Add project root to path for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from config import Config
from strategies import STRATEGIES

# ───────────────────────────────────────────────────────────────
# Expiry horizons to compute targets for
# ───────────────────────────────────────────────────────────────
EXPIRIES = [1, 2, 3, 4, 5, 10]
MIN_IDX = 50           # Skip indicator warmup
FUTURE_MARGIN = 11     # Need 10 future bars for target_10h


def load_data(instrument: str) -> pd.DataFrame:
    """Load cached indicator CSV for one instrument."""
    path = os.path.join(Config.DATA_DIR, f'{instrument}_H1_indicators.csv')
    if not os.path.exists(path):
        raise FileNotFoundError(f'Missing {path}. Run --fetch first.')
    df = pd.read_csv(path, index_col=0, parse_dates=True)
    print(f'  Loaded {instrument}: {len(df)} rows, {df.index[0]} → {df.index[-1]}')
    return df


def compute_signals(df: pd.DataFrame, instrument: str) -> pd.DataFrame:
    """
    Compute strategy signals for all candles and add as columns.
    Returns a new DataFrame with signal/confidence columns (aligned to df index).
    """
    params = Config.STRATEGY_PARAMS
    n = len(df)
    max_idx = n - FUTURE_MARGIN

    # Pre-allocate arrays for each strategy
    sig_cols = {}
    conf_cols = {}
    for name in STRATEGIES:
        sig_cols[name] = np.zeros(n, dtype=int)
        conf_cols[name] = np.zeros(n, dtype=float)

    # Compute signals for each candle
    for idx in range(MIN_IDX, max_idx):
        for name, fn in STRATEGIES.items():
            r = fn(df, params[name], idx)
            if r['signal'] == 'CALL':
                sig_cols[name][idx] = 1
            elif r['signal'] == 'PUT':
                sig_cols[name][idx] = -1
            else:
                sig_cols[name][idx] = 0
            conf_cols[name][idx] = r['confidence']

    # Build result DataFrame indexed like original
    result = pd.DataFrame(index=df.index)
    for name in STRATEGIES:
        result[f'{name}_signal'] = sig_cols[name].astype(int)
        result[f'{name}_confidence'] = conf_cols[name].astype(float)

    return result


def compute_consensus(sig_df: pd.DataFrame) -> tuple:
    """
    Compute weighted consensus direction and total confidence per candle.

    consensus = sign(sum(signal_i * confidence_i))
    total_conf = average confidence of strategies that gave a signal

    Returns (consensus_array, total_conf_array)
    """
    n = len(sig_df)
    consensus = np.zeros(n, dtype=int)
    total_conf = np.zeros(n, dtype=float)

    for idx in range(n):
        weighted = 0.0
        active_count = 0
        for name in STRATEGIES:
            s = sig_df[f'{name}_signal'].iloc[idx]
            c = sig_df[f'{name}_confidence'].iloc[idx]
            if s != 0:
                weighted += s * c
                active_count += 1
        if active_count > 0 and abs(weighted) > 0.01:
            consensus[idx] = 1 if weighted > 0 else -1
            total_conf[idx] = min(abs(weighted) / active_count, 100.0)
        else:
            consensus[idx] = 0
            total_conf[idx] = 0.0

    return consensus, total_conf


def compute_targets(df: pd.DataFrame, sig_df: pd.DataFrame) -> pd.DataFrame:
    """
    Compute target_XYh and best_expiry for each candle using consensus signals.

    For each candle (up to len(df)-FUTURE_MARGIN):
    - Determine consensus direction (weighted by strategy confidences)
    - For each expiry, check if consensus was correct → target_Xh = ±1 or 0
    - best_expiry = expiry with largest price move in consensus direction
    """
    n = len(df)
    max_idx = n - FUTURE_MARGIN

    consensus, total_conf = compute_consensus(sig_df)

    # Pre-allocate target arrays
    targets = {}
    for exp in EXPIRIES:
        targets[f'target_{exp}h'] = np.zeros(n, dtype=int)

    best_expiry_arr = np.zeros(n, dtype=int)

    # Compute targets
    for idx in range(MIN_IDX, max_idx):
        cons = consensus[idx]
        if cons == 0:
            # No consensus → all targets 0, best_expiry 0
            continue

        current_close = df['Close'].iloc[idx]
        if pd.isna(current_close) or current_close == 0:
            continue

        best_move = 0.0
        best_exp = 0

        for exp in EXPIRIES:
            future_close = df['Close'].iloc[idx + exp]
            if pd.isna(future_close):
                continue

            # Check correctness
            if cons == 1:  # CALL
                correct = future_close > current_close
            else:  # PUT
                correct = future_close < current_close

            if correct:
                targets[f'target_{exp}h'][idx] = cons  # 1 or -1
                # Track price move magnitude for best_expiry selection
                move = abs(future_close - current_close) / current_close
                if move > best_move:
                    best_move = move
                    best_exp = exp
            else:
                targets[f'target_{exp}h'][idx] = 0

        best_expiry_arr[idx] = best_exp

    # Build targets DataFrame
    target_df = pd.DataFrame(index=df.index)
    for exp in EXPIRIES:
        target_df[f'target_{exp}h'] = targets[f'target_{exp}h']
    target_df['best_expiry'] = best_expiry_arr

    return target_df


def extract_context_features(df: pd.DataFrame) -> pd.DataFrame:
    """Extract market context features from the indicator DataFrame."""
    ctx = pd.DataFrame(index=df.index)

    ctx['RSI14'] = df['RSI14']
    ctx['ATR14'] = df['ATR14']
    ctx['MACD_hist'] = df['MACD_hist']
    ctx['body_pct'] = df['body_pct']

    # Time-based features
    ctx['hour'] = df.index.hour
    ctx['day_of_week'] = df.index.dayofweek

    return ctx


def build_dataset():
    """Main pipeline: load → compute signals → targets → merge → save."""
    config = Config()
    config.setup_dirs()

    all_rows = []

    for instrument in config.INSTRUMENTS:
        print(f'\n{"=" * 60}')
        print(f'Processing {instrument}...')

        # 1. Load
        df = load_data(instrument)

        # 2. Compute strategy signals
        print(f'  Computing strategy signals for {len(df)} candles...')
        sig_df = compute_signals(df, instrument)

        # 3. Compute targets (consensus-based)
        print(f'  Computing targets for expiries {EXPIRIES}...')
        target_df = compute_targets(df, sig_df)

        # 4. Extract context features
        ctx_df = extract_context_features(df)

        # 5. Merge all into one DataFrame
        merged = sig_df.join(target_df).join(ctx_df)
        merged['instrument'] = instrument
        merged['timestamp'] = df.index

        # 6. Trim to valid range (MIN_IDX to len-FUTURE_MARGIN)
        max_idx = len(df) - FUTURE_MARGIN
        merged = merged.iloc[MIN_IDX:max_idx].copy()

        # 7. Drop rows with any NaN in features/targets
        before_drop = len(merged)
        merged = merged.dropna()
        after_drop = len(merged)
        print(f'  Rows: {before_drop} → {after_drop} after dropna (dropped {before_drop - after_drop})')

        all_rows.append(merged)

    # Combine both instruments
    print(f'\n{"=" * 60}')
    print('Combining instruments...')
    dataset = pd.concat(all_rows, ignore_index=True)

    # Reorder columns for readability
    id_cols = ['instrument', 'timestamp']
    sig_cols = []
    for name in STRATEGIES:
        sig_cols.append(f'{name}_signal')
        sig_cols.append(f'{name}_confidence')
    feature_cols = ['RSI14', 'ATR14', 'MACD_hist', 'body_pct', 'hour', 'day_of_week']
    target_cols = [f'target_{exp}h' for exp in EXPIRIES] + ['best_expiry']
    all_cols = id_cols + sig_cols + feature_cols + target_cols

    # Ensure all columns exist
    missing = [c for c in all_cols if c not in dataset.columns]
    if missing:
        print(f'WARNING: Missing columns: {missing}')
    dataset = dataset[[c for c in all_cols if c in dataset.columns]]

    # Save
    output_path = os.path.join(config.DATA_DIR, 'dataset.csv')
    dataset.to_csv(output_path, index=False)
    print(f'\nDataset saved to: {output_path}')
    print(f'Shape: {dataset.shape[0]} rows × {dataset.shape[1]} columns')

    return dataset


def print_statistics(dataset: pd.DataFrame):
    """Print comprehensive dataset statistics."""
    print(f'\n{"=" * 70}')
    print('DATASET STATISTICS')
    print(f'{"=" * 70}')

    # Instrument distribution
    print(f'\n--- Instrument Distribution ---')
    for instr in ['BITCOIN', 'EURUSD']:
        count = (dataset['instrument'] == instr).sum()
        print(f'  {instr}: {count} rows ({100*count/len(dataset):.1f}%)')

    # Feature columns
    id_cols = ['instrument', 'timestamp']
    sig_cols = [c for c in dataset.columns if c.endswith('_signal')]
    conf_cols = [c for c in dataset.columns if c.endswith('_confidence')]
    feature_cols = ['RSI14', 'ATR14', 'MACD_hist', 'body_pct', 'hour', 'day_of_week']
    target_cols = [c for c in dataset.columns if c.startswith('target_')] + ['best_expiry']

    print(f'\n--- Column Groups ---')
    print(f'  ID columns:        {len(id_cols)} — {id_cols}')
    print(f'  Signal features:   {len(sig_cols)} — {sig_cols}')
    print(f'  Confidence feats:  {len(conf_cols)} — {conf_cols}')
    print(f'  Context features:  {len(feature_cols)} — {feature_cols}')
    print(f'  Target columns:    {len(target_cols)} — {target_cols}')
    print(f'  TOTAL:             {len(dataset.columns)}')

    # NaN check
    nan_counts = dataset.isnull().sum()
    nan_cols = nan_counts[nan_counts > 0]
    if len(nan_cols) > 0:
        print(f'\n--- NaN Summary ---')
        for col, cnt in nan_cols.items():
            print(f'  {col}: {cnt} NaN')
    else:
        print(f'\n--- NaN Check: CLEAN ✓ ---')

    # Target distributions
    print(f'\n--- Target Distribution ---')
    for col in target_cols:
        if col == 'best_expiry':
            # Count per expiry
            print(f'  {col}:')
            for exp in [0] + EXPIRIES:
                cnt = (dataset[col] == exp).sum()
                print(f'    expiry={exp:2d}: {cnt:6d} ({100*cnt/len(dataset):5.1f}%)')
        else:
            vals = dataset[col]
            pos = (vals == 1).sum()
            neg = (vals == -1).sum()
            zero = (vals == 0).sum()
            total = len(vals)
            pct_pos = 100 * pos / total if total > 0 else 0
            pct_neg = 100 * neg / total if total > 0 else 0
            pct_zero = 100 * zero / total if total > 0 else 0
            print(f'  {col}: CALL_won={pos:6d} ({pct_pos:5.1f}%)  '
                  f'PUT_won={neg:6d} ({pct_neg:5.1f}%)  '
                  f'none={zero:6d} ({pct_zero:5.1f}%)')

    # Signal feature stats
    print(f'\n--- Strategy Signal Counts ---')
    for col in sig_cols:
        strat = col.replace('_signal', '')
        calls = (dataset[col] == 1).sum()
        puts = (dataset[col] == -1).sum()
        holds = (dataset[col] == 0).sum()
        print(f'  {strat:20s}: CALL={calls:6d}  PUT={puts:6d}  HOLD={holds:6d}')

    # Context feature stats
    print(f'\n--- Context Feature Stats ---')
    for col in feature_cols:
        s = dataset[col]
        print(f'  {col:15s}: mean={s.mean():.4f}  std={s.std():.4f}  '
              f'min={s.min():.4f}  max={s.max():.4f}')

    # Price movement stats by best_expiry
    print(f'\n--- Consensus Accuracy by Expiry ---')
    for exp in EXPIRIES:
        col = f'target_{exp}h'
        nonzero = (dataset[col] != 0).sum()
        pct = 100 * nonzero / len(dataset)
        print(f'  expiry={exp:2d}h: {nonzero:6d}/{len(dataset)} rows have consensus signal ({pct:.1f}%)')

    # Temporal coverage
    if 'timestamp' in dataset.columns:
        ts = pd.to_datetime(dataset['timestamp'])
        print(f'\n--- Temporal Coverage ---')
        print(f'  From: {ts.min()}')
        print(f'  To:   {ts.max()}')
        print(f'  Span: {ts.max() - ts.min()}')


def main():
    dataset = build_dataset()
    print_statistics(dataset)

    # Quick sanity: verify no look-ahead bias
    print(f'\n--- Sanity Checks ---')
    # Check that signal columns don't depend on future (they're from strategy functions)
    # Check numeric types
    for col in dataset.columns:
        if col in ['instrument', 'timestamp']:
            continue
        if dataset[col].dtype not in [np.int64, np.float64, int, float]:
            print(f'  WARNING: {col} has dtype {dataset[col].dtype}')

    print(f'\nDataset ready for NN training: data/dataset.csv')
    print(f'Features: 5 strategies × 2 (signal+confidence) + 6 context = 16 features')
    print(f'Targets: 6 expiry horizons + 1 best_expiry = 7 targets')
    return dataset


if __name__ == '__main__':
    ds = main()