"""
Training script for Anomaly Detection using Autoencoder.

Note: Ensure TensorFlow version matches production environment for compatibility.
Recommended: tensorflow>=2.13.0,<3.0.0
"""
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from tensorflow import keras
from tensorflow.keras import Model, Input
from tensorflow.keras.layers import Dense
import joblib
import os
import warnings
from preprocessing import engineer_features, get_feature_columns

# Suppress TensorFlow warnings
warnings.filterwarnings('ignore', category=UserWarning)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'


def train_anomaly_detector(
    train_csv_paths: list,
    output_dir: str = 'models',
    encoding_dim: int = 32,
    epochs: int = 50,
    batch_size: int = 256,
    random_state: int = 42
):
    """
    Train autoencoder for anomaly detection.
    
    Args:
        train_csv_paths: List of CSV file paths to load
        output_dir: Directory to save model artifacts
        encoding_dim: Dimension of encoding layer
        epochs: Training epochs
        batch_size: Batch size
        random_state: Random seed
    """
    os.makedirs(output_dir, exist_ok=True)
    np.random.seed(random_state)
    
    # Load and combine data
    print("Loading data...")
    dfs = []
    for csv_path in train_csv_paths:
        if os.path.exists(csv_path):
            df = pd.read_csv(csv_path)
            dfs.append(df)
            print(f"  Loaded {len(df)} rows from {csv_path}")
        else:
            print(f"  Warning: {csv_path} not found, skipping")
    
    if not dfs:
        raise ValueError("No data files found!")
    
    df = pd.concat(dfs, ignore_index=True)
    print(f"Total rows: {len(df)}")
    
    # Engineer features
    print("Engineering features...")
    df = engineer_features(df)
    
    # Get numeric features for autoencoder
    feature_cols = get_feature_columns()
    num_cols = feature_cols['anomaly_numeric']
    
    # Ensure all columns exist
    for col in num_cols:
        if col not in df.columns:
            df[col] = 0.0
    
    X = df[num_cols].copy()
    X = X.fillna(0.0)
    
    print(f"Anomaly features: {num_cols}")
    print(f"Feature matrix shape: {X.shape}")
    
    # Scale features
    print("Scaling features...")
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    # Build autoencoder
    print("Building autoencoder...")
    input_dim = X_scaled.shape[1]
    
    input_layer = Input(shape=(input_dim,))
    encoded = Dense(64, activation='relu')(input_layer)
    encoded = Dense(encoding_dim, activation='relu')(encoded)
    decoded = Dense(64, activation='relu')(encoded)
    decoded = Dense(input_dim, activation='linear')(decoded)
    
    autoencoder = Model(input_layer, decoded)
    autoencoder.compile(optimizer='adam', loss='mse')
    
    print(f"Autoencoder architecture:")
    autoencoder.summary()
    
    # Train
    print("Training autoencoder...")
    history = autoencoder.fit(
        X_scaled, X_scaled,
        epochs=epochs,
        batch_size=batch_size,
        validation_split=0.1,
        verbose=1
    )
    
    # Compute reconstruction errors
    print("Computing reconstruction errors...")
    X_pred = autoencoder.predict(X_scaled, verbose=0)
    recon_errors = np.mean((X_scaled - X_pred) ** 2, axis=1)
    
    df['recon_error'] = recon_errors
    
    # Compute per-job thresholds (97th percentile)
    df['anom_threshold_job'] = df.groupby('job_nm')['recon_error'].transform(
        lambda s: np.percentile(s, 97) if len(s) > 0 else np.percentile(recon_errors, 97)
    )
    df['is_anomaly_job'] = (df['recon_error'] > df['anom_threshold_job']).astype(int)
    
    print(f"\nReconstruction error stats:")
    print(f"  Mean: {recon_errors.mean():.6f}")
    print(f"  Std: {recon_errors.std():.6f}")
    print(f"  Min: {recon_errors.min():.6f}")
    print(f"  Max: {recon_errors.max():.6f}")
    print(f"  97th percentile: {np.percentile(recon_errors, 97):.6f}")
    print(f"  Anomalies detected: {df['is_anomaly_job'].sum()} ({df['is_anomaly_job'].mean()*100:.2f}%)")
    
    # Save model and artifacts
    # Save in Keras format with explicit save_format for compatibility
    model_path = os.path.join(output_dir, 'anomaly_autoencoder_cpu.keras')
    
    print(f"\nSaving model...")
    print(f"TensorFlow version: {keras.__version__ if hasattr(keras, '__version__') else 'unknown'}")
    
    try:
        # Save model without compiling to avoid version-specific config issues
        # This ensures better compatibility across TensorFlow versions
        autoencoder.save(
            model_path,
            save_format='keras',
            include_optimizer=False  # Don't save optimizer for inference-only models
        )
        print(f"✓ Autoencoder saved to {model_path} (Keras format, no optimizer)")
    except TypeError:
        # If include_optimizer is not supported, try without it
        try:
            autoencoder.save(model_path, save_format='keras')
            print(f"✓ Autoencoder saved to {model_path} (Keras format)")
        except Exception as e:
            # Fallback: save without explicit format
            try:
                autoencoder.save(model_path)
                print(f"✓ Autoencoder saved to {model_path} (default format)")
            except Exception as e2:
                print(f"✗ Error saving model: {e2}")
                raise
    except Exception as e:
        print(f"✗ Error saving model: {e}")
        raise
    
    scaler_path = os.path.join(output_dir, 'anomaly_scaler.joblib')
    joblib.dump(scaler, scaler_path)
    print(f"Scaler saved to {scaler_path}")
    
    feature_list_path = os.path.join(output_dir, 'anomaly_features.joblib')
    joblib.dump(num_cols, feature_list_path)
    print(f"Feature list saved to {feature_list_path}")
    
    # Save global threshold (97th percentile)
    global_threshold = np.percentile(recon_errors, 97)
    threshold_path = os.path.join(output_dir, 'anomaly_threshold.joblib')
    joblib.dump(global_threshold, threshold_path)
    print(f"Global threshold saved to {threshold_path} (value: {global_threshold:.6f})")
    
    return autoencoder, scaler, X_scaled, recon_errors


if __name__ == '__main__':
    import sys
    import glob
    
    # Default CSV paths (can be overridden via command line)
    if len(sys.argv) > 1:
        # Handle glob patterns (works in both PowerShell and bash)
        csv_paths = []
        for arg in sys.argv[1:]:
            # Expand glob patterns
            expanded = glob.glob(arg)
            if expanded:
                csv_paths.extend(expanded)
            else:
                # If no match, use as-is (might be a specific file)
                csv_paths.append(arg)
    else:
        # Default: all CSV files in data directory
        csv_paths = glob.glob('data/*.csv')
        if not csv_paths:
            csv_paths = [
                'data/true_export_report_20260120.csv',
                'data/true_export_report_20260121.csv'
            ]
    
    if not csv_paths:
        print("Error: No CSV files found!")
        print("Usage: python train_anomaly.py [file1.csv] [file2.csv] ...")
        print("   or: python train_anomaly.py data/*.csv")
        sys.exit(1)
    
    print(f"Training with {len(csv_paths)} file(s):")
    for path in csv_paths:
        print(f"  - {path}")
    print()
    
    train_anomaly_detector(csv_paths)