"""
Comprehensive data augmentation and model improvement pipeline.

Augmentation strategies:
1. Variable subsampling: randomly drop variables to create new graph topologies
2. Sample-size variation: subsample rows from existing large-N datasets
3. Noise injection: add random noise to some variables

Then trains multiple model architectures and does a full comparison.
"""
import os
import sys
import numpy as np
import pandas as pd
import json
import logging
import warnings
import time
from itertools import combinations

warnings.filterwarnings('ignore')
logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')
logger = logging.getLogger(__name__)
logging.getLogger('causallearn').setLevel(logging.ERROR)

sys.path.insert(0, '/app')
from causal_selection.data.generator import (
    load_bn_model, get_true_dag_adjmat, dag_to_cpdag, sample_dataset,
    ALL_NETWORKS, MEDIUM_NETWORKS, LARGE_NETWORKS, get_network_tier
)
from causal_selection.discovery.algorithms import run_algorithm, ALGORITHM_POOL
from causal_selection.discovery.evaluator import evaluate_algorithm_result
from causal_selection.features.extractor import extract_all_features, FEATURE_NAMES
from causal_selection.meta_learner.trainer import (
    load_meta_dataset, evaluate_lono_cv, train_meta_learner,
    save_model, get_feature_importance, ALGO_NAMES, RESULTS_DIR
)

from sklearn.ensemble import (
    RandomForestRegressor, GradientBoostingRegressor,
    RandomForestClassifier, GradientBoostingClassifier
)
from sklearn.multioutput import MultiOutputRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error
import joblib


# ==============================================================
# AUGMENTATION
# ==============================================================

def augment_all(networks_for_varsub=None, n_varsub=3, drop_frac=0.3,
                networks_for_samplesub=None, n_samplesub=2):
    """Run all augmentation strategies and return combined augmented data."""
    
    all_feats, all_shds, all_nshds, all_cfgs = [], [], [], []
    
    # Strategy 1: Variable subsampling
    logger.info("="*60)
    logger.info("AUGMENTATION: Variable Subsampling")
    logger.info("="*60)
    
    if networks_for_varsub is None:
        # Only networks with >8 variables
        networks_for_varsub = ['sachs', 'alarm', 'child', 'insurance', 
                                'water', 'barley', 'mildew',
                                'hailfinder', 'hepar2']
    
    for net_name in networks_for_varsub:
        try:
            model = load_bn_model(net_name)
            true_dag, node_names = get_true_dag_adjmat(model)
            n_vars = len(node_names)
            
            if n_vars < 8:
                continue
            
            tier = get_network_tier(net_name)
            timeout = {'small': 60, 'medium': 90, 'large': 120}[tier]
            
            for aug_i in range(n_varsub):
                rng = np.random.RandomState(200 + aug_i * 100 + hash(net_name) % 100)
                
                # Keep 60-80% of variables
                keep_frac = 1.0 - drop_frac + rng.uniform(-0.1, 0.1)
                keep_frac = max(0.5, min(0.85, keep_frac))
                n_keep = max(5, int(n_vars * keep_frac))
                keep_idx = sorted(rng.choice(n_vars, n_keep, replace=False))
                
                sub_dag = true_dag[np.ix_(keep_idx, keep_idx)]
                sub_cpdag = dag_to_cpdag(sub_dag)
                sub_names = [node_names[i] for i in keep_idx]
                
                n_samples = rng.choice([500, 1000, 2000])
                df_full = sample_dataset(model, n_samples, seed=200 + aug_i)
                df_sub = df_full[sub_names].copy()
                df_sub.columns = [f'X{i}' for i in range(len(sub_names))]
                
                logger.info(f"  VarSub {net_name} #{aug_i}: {n_vars}->{n_keep} vars, N={n_samples}")
                
                f, s, ns, c = _run_single(df_sub, sub_cpdag, 
                                           f'{net_name}_vs{aug_i}', n_samples, 
                                           200+aug_i, n_keep, timeout)
                if f is not None:
                    all_feats.append(f)
                    all_shds.append(s)
                    all_nshds.append(ns)
                    all_cfgs.append(c)
                    
        except Exception as e:
            logger.error(f"VarSub failed for {net_name}: {e}")
    
    # Strategy 2: Sample-size subsampling from existing large-N datasets
    logger.info("\n" + "="*60)
    logger.info("AUGMENTATION: Sample Size Variation")
    logger.info("="*60)
    
    if networks_for_samplesub is None:
        networks_for_samplesub = ['asia', 'cancer', 'earthquake', 'sachs', 
                                   'survey', 'alarm', 'child']
    
    sub_sample_sizes = [300, 750, 1500, 3000]
    
    for net_name in networks_for_samplesub:
        try:
            model = load_bn_model(net_name)
            true_dag, node_names = get_true_dag_adjmat(model)
            true_cpdag = dag_to_cpdag(true_dag)
            n_vars = len(node_names)
            tier = get_network_tier(net_name)
            timeout = {'small': 60, 'medium': 90, 'large': 120}[tier]
            
            for ss_i, n_samples in enumerate(sub_sample_sizes):
                seed = 300 + ss_i
                df = sample_dataset(model, n_samples, seed=seed)
                
                logger.info(f"  SampleSub {net_name} N={n_samples} seed={seed}")
                
                f, s, ns, c = _run_single(df, true_cpdag, 
                                           f'{net_name}_ss{ss_i}', n_samples,
                                           seed, n_vars, timeout)
                if f is not None:
                    all_feats.append(f)
                    all_shds.append(s)
                    all_nshds.append(ns)
                    all_cfgs.append(c)
                    
        except Exception as e:
            logger.error(f"SampleSub failed for {net_name}: {e}")
    
    # Strategy 3: Noise injection on small networks
    logger.info("\n" + "="*60)
    logger.info("AUGMENTATION: Noise Injection")
    logger.info("="*60)
    
    noise_networks = ['asia', 'sachs', 'survey', 'cancer', 'earthquake']
    
    for net_name in noise_networks:
        try:
            model = load_bn_model(net_name)
            true_dag, node_names = get_true_dag_adjmat(model)
            true_cpdag = dag_to_cpdag(true_dag)
            n_vars = len(node_names)
            timeout = 60
            
            for noise_i, noise_frac in enumerate([0.05, 0.10]):
                seed = 400 + noise_i
                n_samples = 1000
                df = sample_dataset(model, n_samples, seed=seed)
                
                # Inject random category flips
                rng = np.random.RandomState(seed)
                n_flip = int(n_samples * n_vars * noise_frac)
                for _ in range(n_flip):
                    r = rng.randint(n_samples)
                    c = rng.randint(n_vars)
                    max_val = df.iloc[:, c].max()
                    df.iloc[r, c] = rng.randint(0, max_val + 1)
                
                logger.info(f"  Noise {net_name} frac={noise_frac}")
                
                f, s, ns, c = _run_single(df, true_cpdag,
                                           f'{net_name}_n{noise_i}', n_samples,
                                           seed, n_vars, timeout)
                if f is not None:
                    all_feats.append(f)
                    all_shds.append(s)
                    all_nshds.append(ns)
                    all_cfgs.append(c)
                    
        except Exception as e:
            logger.error(f"Noise failed for {net_name}: {e}")
    
    return all_feats, all_shds, all_nshds, all_cfgs


def _run_single(df, true_cpdag, net_label, n_samples, seed, n_vars, timeout):
    """Run feature extraction + all algorithms on one config."""
    try:
        features = extract_all_features(df, n_probe_triplets=60)
        
        shd_row = {}
        nshd_row = {}
        max_possible = n_vars * (n_vars - 1) // 2
        
        for algo_name in ALGO_NAMES:
            result = run_algorithm(algo_name, df, timeout_sec=timeout)
            metrics = evaluate_algorithm_result(result, true_cpdag)
            shd_row[algo_name] = metrics['shd']
            nshd_row[algo_name] = metrics['normalized_shd']
        
        feat_row = {name: features.get(name, 0.0) for name in FEATURE_NAMES}
        config = {
            'network': net_label,
            'n_samples': n_samples,
            'seed': seed,
            'n_variables': n_vars,
            'n_true_edges': int(((true_cpdag + true_cpdag.T) > 0).sum() // 2),
        }
        
        # Log best algo
        best = min(shd_row, key=shd_row.get)
        logger.info(f"    Best: {best} SHD={shd_row[best]}")
        
        return feat_row, shd_row, nshd_row, config
        
    except Exception as e:
        logger.error(f"    Failed: {e}")
        return None, None, None, None


# ==============================================================
# PAIRWISE RANKING MODEL
# ==============================================================

def train_pairwise_ranking(X, Y_nshd, configs):
    """Train pairwise ranking classifiers: for each (algo_i, algo_j) pair,
    train a classifier to predict whether algo_i beats algo_j.
    
    At inference: count wins for each algorithm, rank by win count.
    """
    n_algos = len(ALGO_NAMES)
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    pair_models = {}
    pair_accuracies = {}
    
    for i in range(n_algos):
        for j in range(i+1, n_algos):
            algo_i, algo_j = ALGO_NAMES[i], ALGO_NAMES[j]
            
            # Label: 1 if algo_i has lower nSHD (better) than algo_j
            y = (Y_nshd.iloc[:, i] < Y_nshd.iloc[:, j]).astype(int).values
            
            # Skip if one always wins
            if y.mean() == 0 or y.mean() == 1:
                pair_models[(i,j)] = None
                pair_accuracies[(i,j)] = y.mean()
                continue
            
            clf = GradientBoostingClassifier(
                n_estimators=200, max_depth=3, learning_rate=0.05,
                random_state=42
            )
            clf.fit(X_scaled, y)
            
            train_acc = clf.score(X_scaled, y)
            pair_models[(i,j)] = clf
            pair_accuracies[(i,j)] = train_acc
    
    return pair_models, scaler, pair_accuracies


def predict_pairwise_ranking(pair_models, scaler, X_new, k=3):
    """Use pairwise models to rank algorithms via win-count."""
    X_scaled = scaler.transform(X_new)
    n_algos = len(ALGO_NAMES)
    n_samples = X_scaled.shape[0]
    
    results = []
    for idx in range(n_samples):
        wins = np.zeros(n_algos)
        x = X_scaled[idx:idx+1]
        
        for i in range(n_algos):
            for j in range(i+1, n_algos):
                model = pair_models.get((i,j))
                if model is None:
                    continue
                pred = model.predict(x)[0]
                if pred == 1:  # algo_i wins
                    wins[i] += 1
                else:
                    wins[j] += 1
        
        ranking = np.argsort(-wins)  # most wins first
        results.append(ranking[:k])
    
    return np.array(results)


def evaluate_pairwise_lono(X, Y_nshd, configs, k=3):
    """LONO-CV for pairwise ranking model."""
    networks = configs['network'].values
    unique_nets = sorted(configs['network'].unique())
    # For augmented data, group by base network name
    base_nets = configs['network'].apply(lambda x: x.split('_')[0]).values
    unique_base = sorted(set(base_nets))
    
    top_k_hits = 0
    regrets = []
    total = 0
    
    for test_base in unique_base:
        test_mask = base_nets == test_base
        train_mask = ~test_mask
        
        if train_mask.sum() < 5 or test_mask.sum() == 0:
            continue
        
        X_train = X.values[train_mask]
        Y_train = Y_nshd[train_mask]
        X_test = X.values[test_mask]
        Y_test = Y_nshd.values[test_mask]
        
        # Train pairwise models
        scaler = StandardScaler()
        X_train_s = scaler.fit_transform(X_train)
        
        n_algos = len(ALGO_NAMES)
        pair_models = {}
        
        for i in range(n_algos):
            for j in range(i+1, n_algos):
                y = (Y_train.iloc[:, i] < Y_train.iloc[:, j]).astype(int).values
                if y.mean() == 0 or y.mean() == 1:
                    pair_models[(i,j)] = None
                    continue
                clf = GradientBoostingClassifier(
                    n_estimators=100, max_depth=3, learning_rate=0.05,
                    random_state=42
                )
                clf.fit(X_train_s, y)
                pair_models[(i,j)] = clf
        
        # Predict
        X_test_s = scaler.transform(X_test)
        
        for idx in range(len(X_test_s)):
            wins = np.zeros(n_algos)
            x = X_test_s[idx:idx+1]
            
            for i in range(n_algos):
                for j in range(i+1, n_algos):
                    m = pair_models.get((i,j))
                    if m is None:
                        continue
                    if m.predict(x)[0] == 1:
                        wins[i] += 1
                    else:
                        wins[j] += 1
            
            pred_top_k = np.argsort(-wins)[:k]
            true_best = np.argmin(Y_test[idx])
            
            if true_best in pred_top_k:
                top_k_hits += 1
            
            oracle = Y_test[idx, true_best]
            selected = min(Y_test[idx, a] for a in pred_top_k)
            regrets.append(selected - oracle)
            total += 1
    
    hit_rate = top_k_hits / total if total > 0 else 0
    mean_regret = np.mean(regrets) if regrets else 0
    
    return {
        'top_k_hit_rate': hit_rate,
        'mean_regret': mean_regret,
        'median_regret': np.median(regrets) if regrets else 0,
        'n_evaluated': total,
    }


# ==============================================================
# MAIN
# ==============================================================

if __name__ == '__main__':
    start_time = time.time()
    
    # Step 1: Augment
    print("="*80)
    print("STEP 1: DATA AUGMENTATION")
    print("="*80)
    
    feats, shds, nshds, cfgs = augment_all(
        n_varsub=2, drop_frac=0.3,
        n_samplesub=2,
    )
    
    print(f"\nGenerated {len(cfgs)} augmented configs")
    
    # Merge with original
    X_orig, Y_shd_orig, Y_nshd_orig, configs_orig = load_meta_dataset()
    
    X_aug = pd.DataFrame(feats, columns=FEATURE_NAMES)
    Y_shd_aug = pd.DataFrame(shds, columns=ALGO_NAMES)
    Y_nshd_aug = pd.DataFrame(nshds, columns=ALGO_NAMES)
    configs_aug = pd.DataFrame(cfgs)
    
    X_all = pd.concat([X_orig, X_aug], ignore_index=True)
    Y_shd_all = pd.concat([Y_shd_orig, Y_shd_aug], ignore_index=True)
    Y_nshd_all = pd.concat([Y_nshd_orig, Y_nshd_aug], ignore_index=True)
    configs_all = pd.concat([configs_orig, configs_aug], ignore_index=True)
    
    print(f"Total dataset: {len(configs_all)} configs "
          f"({len(configs_orig)} original + {len(configs_aug)} augmented)")
    
    # Save augmented data
    X_all.to_csv(os.path.join(RESULTS_DIR, 'meta_features.csv'), index=False)
    Y_shd_all.to_csv(os.path.join(RESULTS_DIR, 'shd_matrix.csv'), index=False)
    Y_nshd_all.to_csv(os.path.join(RESULTS_DIR, 'normalized_shd_matrix.csv'), index=False)
    configs_all.to_csv(os.path.join(RESULTS_DIR, 'configs.csv'), index=False)
    
    # Step 2: Model comparison
    print("\n" + "="*80)
    print("STEP 2: MODEL COMPARISON (LONO-CV)")
    print("="*80)
    
    # Reload augmented data
    X, Y_shd, Y_nshd, configs = load_meta_dataset()
    
    print(f"\n{'Model':25s} {'Top3Hit':>8s} {'NDCG@3':>8s} {'Regret':>8s}")
    print("-"*55)
    
    model_configs = [
        ('RF-200', 'rf', {'n_estimators': 200}),
        ('RF-500', 'rf', {'n_estimators': 500}),
        ('GBM-500-lr05', 'gbm', {'n_estimators': 500, 'max_depth': 3, 'learning_rate': 0.05}),
        ('GBM-300-lr01', 'gbm', {'n_estimators': 300, 'max_depth': 4, 'learning_rate': 0.01}),
        ('GBM-200-lr1', 'gbm', {'n_estimators': 200, 'max_depth': 5, 'learning_rate': 0.1}),
    ]
    
    best_hit = 0
    best_config = None
    
    for name, mtype, kwargs in model_configs:
        r = evaluate_lono_cv(X, Y_nshd, configs, model_type=mtype, k=3, **kwargs)
        o = r['overall']
        print(f"{name:25s} {o['top_k_hit_rate']:8.3f} {o['ndcg_at_k']:8.3f} {o['mean_regret']:8.4f}")
        if o['top_k_hit_rate'] > best_hit:
            best_hit = o['top_k_hit_rate']
            best_config = (name, mtype, kwargs, o)
    
    # Pairwise ranking
    print(f"\n{'Pairwise-GBM':25s}", end="")
    pw_results = evaluate_pairwise_lono(X, Y_nshd, configs, k=3)
    print(f" {pw_results['top_k_hit_rate']:8.3f} {'N/A':>8s} {pw_results['mean_regret']:8.4f}")
    
    if pw_results['top_k_hit_rate'] > best_hit:
        best_hit = pw_results['top_k_hit_rate']
        best_config = ('Pairwise-GBM', 'pairwise', {}, pw_results)
    
    print(f"\n{'='*55}")
    print(f"BEST MODEL: {best_config[0]} (hit rate={best_hit:.3f})")
    print(f"{'='*55}")
    
    # Train & save best multi-output model
    if best_config[1] != 'pairwise':
        model, scaler = train_meta_learner(X, Y_nshd, 
                                            model_type=best_config[1], 
                                            **best_config[2])
        save_model(model, scaler)
        
        avg_imp, _ = get_feature_importance(model)
        print("\nTop 10 Features:")
        for feat, imp in sorted(avg_imp.items(), key=lambda x: -x[1])[:10]:
            print(f"  {feat:30s}: {imp:.4f}")
    else:
        # Save pairwise model separately
        print("Pairwise model is best - training final version...")
        pair_models, scaler, _ = train_pairwise_ranking(X, Y_nshd, configs)
        os.makedirs('/app/causal_selection/models', exist_ok=True)
        joblib.dump({'pair_models': pair_models, 'scaler': scaler}, 
                    '/app/causal_selection/models/pairwise_model.pkl')
        # Also train and save best multi-output as fallback
        best_mo = [c for c in model_configs if c[0] != 'Pairwise-GBM']
        best_mo_hit = 0
        best_mo_cfg = model_configs[0]
        for name, mtype, kwargs in model_configs:
            r = evaluate_lono_cv(X, Y_nshd, configs, model_type=mtype, k=3, **kwargs)
            if r['overall']['top_k_hit_rate'] > best_mo_hit:
                best_mo_hit = r['overall']['top_k_hit_rate']
                best_mo_cfg = (name, mtype, kwargs)
        model, scaler = train_meta_learner(X, Y_nshd, model_type=best_mo_cfg[1], **best_mo_cfg[2])
        save_model(model, scaler)
    
    elapsed = time.time() - start_time
    print(f"\nTotal time: {elapsed/60:.1f} minutes")
    
    # Save summary
    summary = {
        'n_configs_original': int(len(configs_orig)),
        'n_configs_augmented': int(len(configs_aug)),
        'n_configs_total': int(len(configs_all)),
        'best_model': best_config[0],
        'best_top3_hit_rate': float(best_hit),
        'best_metrics': {k: float(v) if isinstance(v, (float, np.floating)) else v 
                        for k, v in best_config[3].items()},
    }
    with open(os.path.join(RESULTS_DIR, 'improvement_summary.json'), 'w') as f:
        json.dump(summary, f, indent=2)
    
    print(f"\nSummary saved to {RESULTS_DIR}/improvement_summary.json")