"""
SOLUȚIA INTENDED - Machine Learning "by the book"

Pipeline complet:
1. Încărcare date
2. Exploratory Data Analysis
3. Preprocessing (missing values, encoding)
4. Feature Engineering
5. Model Selection & Training
6. Hyperparameter Tuning
7. Predicții finale
"""

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split, cross_val_score, GridSearchCV
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.impute import SimpleImputer
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score, classification_report
import warnings
warnings.filterwarnings('ignore')

# ============================================================
# 1. ÎNCĂRCARE DATE
# ============================================================
print("=" * 60)
print("1. ÎNCĂRCARE DATE")
print("=" * 60)

train_df = pd.read_csv('train.csv')
test_df = pd.read_csv('test.csv')

print(f"Train shape: {train_df.shape}")
print(f"Test shape: {test_df.shape}")
print(f"\nColoane: {list(train_df.columns)}")

# ============================================================
# 2. EXPLORATORY DATA ANALYSIS
# ============================================================
print("\n" + "=" * 60)
print("2. EXPLORATORY DATA ANALYSIS")
print("=" * 60)

print("\n--- Info despre date ---")
print(train_df.info())

print("\n--- Missing values ---")
print(train_df.isnull().sum())

print("\n--- Distribuția target ---")
print(train_df['PriorityLevel'].value_counts())

print("\n--- Statistici numerice ---")
print(train_df.describe())

print("\n--- Valori unice pentru categorice ---")
categorical_cols = ['DisasterType', 'Region', 'InfrastructureDamage', 'WeatherCondition']
for col in categorical_cols:
    print(f"{col}: {train_df[col].unique()}")

# ============================================================
# 3. PREPROCESSING
# ============================================================
print("\n" + "=" * 60)
print("3. PREPROCESSING")
print("=" * 60)

def preprocess_data(train_df, test_df):
    """
    Preprocessing complet:
    - Handle missing values
    - Encode categorical variables
    - Feature scaling (opțional pentru tree-based models)
    """
    train = train_df.copy()
    test = test_df.copy()
    
    # Salvăm SampleID și target
    train_ids = train['SampleID']
    test_ids = test['SampleID']
    target = train['PriorityLevel']
    
    # Drop SampleID și target din features
    train = train.drop(columns=['SampleID', 'PriorityLevel'])
    test = test.drop(columns=['SampleID'])
    
    # --- 3.1 Handle Missing Values ---
    # Numeric columns
    numeric_cols = ['Month', 'Casualties', 'AffectedAreaKm2', 'PopulationDensity', 
                    'ResponseTimeHours', 'EmergencyTeamsAvailable']
    
    # Categorical columns
    categorical_cols = ['DisasterType', 'Region', 'InfrastructureDamage', 'WeatherCondition']
    
    # Impute numeric cu median
    num_imputer = SimpleImputer(strategy='median')
    train[numeric_cols] = num_imputer.fit_transform(train[numeric_cols])
    test[numeric_cols] = num_imputer.transform(test[numeric_cols])
    
    # Impute categorical cu mode
    cat_imputer = SimpleImputer(strategy='most_frequent')
    train[categorical_cols] = cat_imputer.fit_transform(train[categorical_cols])
    test[categorical_cols] = cat_imputer.transform(test[categorical_cols])
    
    print("Missing values handled ✓")
    
    # --- 3.2 Encode Categorical Variables ---
    label_encoders = {}
    
    for col in categorical_cols:
        le = LabelEncoder()
        # Fit pe toate valorile posibile din train + test
        all_values = pd.concat([train[col], test[col]]).unique()
        le.fit(all_values)
        
        train[col + '_encoded'] = le.transform(train[col])
        test[col + '_encoded'] = le.transform(test[col])
        label_encoders[col] = le
    
    print("Categorical encoding done ✓")
    
    # --- 3.3 Encode Target ---
    target_encoder = LabelEncoder()
    target_encoded = target_encoder.fit_transform(target)
    print(f"Target classes: {list(target_encoder.classes_)}")
    
    return train, test, target_encoded, target_encoder, train_ids, test_ids, label_encoders

train_processed, test_processed, y_encoded, target_encoder, train_ids, test_ids, label_encoders = \
    preprocess_data(train_df, test_df)

# ============================================================
# 4. FEATURE ENGINEERING
# ============================================================
print("\n" + "=" * 60)
print("4. FEATURE ENGINEERING")
print("=" * 60)

def create_features(df):
    """
    Creează features noi bazate pe domain knowledge.
    """
    df = df.copy()
    
    # Impact Score = Area * Population Density
    df['ImpactScore'] = df['AffectedAreaKm2'] * df['PopulationDensity'] / 1000
    
    # Log transforms pentru distribuții skewed
    df['Log_Casualties'] = np.log1p(df['Casualties'])
    df['Log_Area'] = np.log1p(df['AffectedAreaKm2'])
    df['Log_Impact'] = np.log1p(df['ImpactScore'])
    
    # Binned casualties
    df['Casualties_Bin'] = pd.cut(df['Casualties'], 
                                   bins=[-1, 0, 4, 19, 49, 1000], 
                                   labels=[0, 1, 2, 3, 4]).astype(int)
    
    # Response time category
    df['ResponseTime_Bin'] = pd.cut(df['ResponseTimeHours'],
                                     bins=[-1, 2, 12, 24, 100],
                                     labels=[0, 1, 2, 3]).astype(int)
    
    # Impact category
    df['Impact_Bin'] = pd.cut(df['ImpactScore'],
                               bins=[-1, 20, 100, 500, 10000],
                               labels=[0, 1, 2, 3]).astype(int)
    
    # Is urban/coastal (high density areas)
    df['IsHighDensity'] = (df['PopulationDensity'] > 500).astype(int)
    
    # Severe weather flag
    df['IsSevereWeather'] = df['WeatherCondition_encoded'].isin([2, 3]).astype(int)  # Severe or Extreme
    
    # Critical infrastructure flag
    df['IsCriticalInfra'] = (df['InfrastructureDamage_encoded'] >= 2).astype(int)  # Severe or Critical
    
    return df

train_featured = create_features(train_processed)
test_featured = create_features(test_processed)

print(f"Features create: {train_featured.shape[1]} total")
print(f"New features: ImpactScore, Log transforms, Bins, Flags")

# ============================================================
# 5. SELECT FEATURES
# ============================================================
print("\n" + "=" * 60)
print("5. SELECT FEATURES")
print("=" * 60)

# Selectăm feature-urile pentru model
feature_columns = [
    # Original numeric
    'Casualties', 'AffectedAreaKm2', 'PopulationDensity', 
    'ResponseTimeHours', 'EmergencyTeamsAvailable', 'Month',
    # Encoded categorical
    'DisasterType_encoded', 'Region_encoded', 
    'InfrastructureDamage_encoded', 'WeatherCondition_encoded',
    # Engineered features
    'ImpactScore', 'Log_Casualties', 'Log_Area', 'Log_Impact',
    'Casualties_Bin', 'ResponseTime_Bin', 'Impact_Bin',
    'IsHighDensity', 'IsSevereWeather', 'IsCriticalInfra'
]

X_train = train_featured[feature_columns]
X_test = test_featured[feature_columns]
y_train = y_encoded

print(f"X_train shape: {X_train.shape}")
print(f"X_test shape: {X_test.shape}")
print(f"Features used: {len(feature_columns)}")

# ============================================================
# 6. MODEL SELECTION & BASELINE
# ============================================================
print("\n" + "=" * 60)
print("6. MODEL SELECTION & BASELINE")
print("=" * 60)

# Split pentru validare locală
X_tr, X_val, y_tr, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42, stratify=y_train)

# Test mai multe modele
models = {
    'DecisionTree': DecisionTreeClassifier(random_state=42),
    'RandomForest': RandomForestClassifier(n_estimators=100, random_state=42),
    'GradientBoosting': GradientBoostingClassifier(n_estimators=100, random_state=42)
}

print("\n--- Baseline Results (CV=5) ---")
for name, model in models.items():
    scores = cross_val_score(model, X_train, y_train, cv=5, scoring='accuracy')
    print(f"{name}: {scores.mean():.4f} (+/- {scores.std()*2:.4f})")

# ============================================================
# 7. HYPERPARAMETER TUNING
# ============================================================
print("\n" + "=" * 60)
print("7. HYPERPARAMETER TUNING (GridSearchCV)")
print("=" * 60)

# Random Forest tuning
rf_params = {
    'n_estimators': [100, 200, 300],
    'max_depth': [10, 20, 30, None],
    'min_samples_split': [2, 5, 10],
    'min_samples_leaf': [1, 2, 4]
}

print("\nTuning Random Forest...")
rf_grid = GridSearchCV(
    RandomForestClassifier(random_state=42),
    rf_params,
    cv=5,
    scoring='accuracy',
    n_jobs=-1,
    verbose=0
)
rf_grid.fit(X_train, y_train)

print(f"Best RF params: {rf_grid.best_params_}")
print(f"Best RF CV score: {rf_grid.best_score_:.4f}")

# Decision Tree tuning (poate atinge 100% dacă regula e descoperibilă)
dt_params = {
    'max_depth': [10, 15, 20, 25, 30, None],
    'min_samples_split': [2, 3, 5],
    'min_samples_leaf': [1, 2],
    'criterion': ['gini', 'entropy']
}

print("\nTuning Decision Tree...")
dt_grid = GridSearchCV(
    DecisionTreeClassifier(random_state=42),
    dt_params,
    cv=5,
    scoring='accuracy',
    n_jobs=-1,
    verbose=0
)
dt_grid.fit(X_train, y_train)

print(f"Best DT params: {dt_grid.best_params_}")
print(f"Best DT CV score: {dt_grid.best_score_:.4f}")

# ============================================================
# 8. FINAL MODEL TRAINING
# ============================================================
print("\n" + "=" * 60)
print("8. FINAL MODEL TRAINING")
print("=" * 60)

# Alegem cel mai bun model
best_rf = rf_grid.best_estimator_
best_dt = dt_grid.best_estimator_

# Antrenăm pe toate datele de train
best_rf.fit(X_train, y_train)
best_dt.fit(X_train, y_train)

# Evaluare pe validation split
y_pred_rf = best_rf.predict(X_val)
y_pred_dt = best_dt.predict(X_val)

print(f"\nValidation Accuracy:")
print(f"  Random Forest: {accuracy_score(y_val, y_pred_rf):.4f}")
print(f"  Decision Tree: {accuracy_score(y_val, y_pred_dt):.4f}")

# Train accuracy (pentru a vedea dacă modelul a învățat regula)
y_train_pred_rf = best_rf.predict(X_train)
y_train_pred_dt = best_dt.predict(X_train)

print(f"\nTrain Accuracy:")
print(f"  Random Forest: {accuracy_score(y_train, y_train_pred_rf):.4f}")
print(f"  Decision Tree: {accuracy_score(y_train, y_train_pred_dt):.4f}")

# Feature importance
print("\n--- Feature Importance (Random Forest) ---")
importance_df = pd.DataFrame({
    'Feature': feature_columns,
    'Importance': best_rf.feature_importances_
}).sort_values('Importance', ascending=False)
print(importance_df.head(10).to_string(index=False))

# ============================================================
# 9. PREDICȚII PE TEST
# ============================================================
print("\n" + "=" * 60)
print("9. PREDICȚII PE TEST")
print("=" * 60)

# Folosim modelul cu train accuracy 100% (Decision Tree)
if accuracy_score(y_train, y_train_pred_dt) > accuracy_score(y_train, y_train_pred_rf):
    final_model = best_dt
    model_name = "Decision Tree"
else:
    final_model = best_rf
    model_name = "Random Forest"

print(f"Using: {model_name}")

# Predicții
y_test_pred_encoded = final_model.predict(X_test)
y_test_pred = target_encoder.inverse_transform(y_test_pred_encoded)

print(f"Predictions made: {len(y_test_pred)}")
print(f"Distribution: {pd.Series(y_test_pred).value_counts().to_dict()}")

# ============================================================
# 10. GENERARE SUBMISSION
# ============================================================
print("\n" + "=" * 60)
print("10. GENERARE SUBMISSION")
print("=" * 60)

# Subtask 1: Câte Hurricane în test?
hurricane_count = len(test_df[test_df['DisasterType'] == 'Hurricane'])

# Subtask 2: Cel mai frecvent InfrastructureDamage în test?
most_common_infra = test_df['InfrastructureDamage'].mode()[0]

# Creăm submission
submission_rows = []

# Subtask 1
submission_rows.append({
    'subtaskID': 1,
    'datapointID': 1,
    'answer': hurricane_count
})

# Subtask 2
submission_rows.append({
    'subtaskID': 2,
    'datapointID': 2,
    'answer': most_common_infra
})

# Subtask 3 - predicții
for idx, sample_id in enumerate(test_ids):
    submission_rows.append({
        'subtaskID': 3,
        'datapointID': sample_id,
        'answer': y_test_pred[idx]
    })

submission_df = pd.DataFrame(submission_rows)
submission_df.to_csv('submission.csv', index=False)

print(f"\n=== SUBMISSION GENERATED ===")
print(f"Subtask 1 (Hurricane count): {hurricane_count}")
print(f"Subtask 2 (Most common infra): {most_common_infra}")
print(f"Subtask 3 (Predictions): {len(y_test_pred)} rows")
print(f"\nTotal rows: {len(submission_df)}")
print(f"\nSaved to: submission.csv")

print("\n--- First 10 rows ---")
print(submission_df.head(10).to_string(index=False))

# ============================================================
# 11. VERIFICARE (doar dacă avem validation.csv)
# ============================================================
print("\n" + "=" * 60)
print("11. VERIFICARE FINALĂ")
print("=" * 60)

try:
    validation_df = pd.read_csv('validation.csv')
    
    # Merge pentru comparație
    sub3 = submission_df[submission_df['subtaskID'] == 3].copy()
    merged = sub3.merge(validation_df[['SampleID', 'PriorityLevel']], 
                        left_on='datapointID', right_on='SampleID')
    
    final_accuracy = (merged['answer'] == merged['PriorityLevel']).mean()
    
    print(f"\n✓ FINAL ACCURACY: {final_accuracy * 100:.2f}%")
    print(f"  Correct: {(merged['answer'] == merged['PriorityLevel']).sum()} / {len(merged)}")
    
    if final_accuracy == 1.0:
        print("\n🎉 PERFECT SCORE! 100% ACCURACY ACHIEVED!")
    
except FileNotFoundError:
    print("validation.csv not found (normal pentru participanți)")
