import numpy as np
import pandas as pd
import os
import gc
import warnings
from pathlib import Path
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
import subprocess
from datetime import datetime

warnings.filterwarnings('ignore')

# Try to import xgboost/lightgbm, install if missing
try:
    import xgboost as xgb
except Exception:
    subprocess.run(["pip", "install", "-q", "xgboost"], check=False)
    import xgboost as xgb

try:
    import lightgbm as lgb
except Exception:
    subprocess.run(["pip", "install", "-q", "lightgbm"], check=False)
    import lightgbm as lgb

class CloudMLEOptimizer:
    def __init__(self):
        print("ARIEL SPACE CHALLENGE - MLE-STAR CLOUD EDITION (v3)")
        print(f"Started at: {datetime.now()}")

class AdvancedSpaceDataProcessor:
    def load_data_efficiently(self, competition_path="/kaggle/input/ariel-data-challenge-2025"):
        print("\nLOADING ARIEL SPACE TELESCOPE DATA:")
        try:
            train_path = f"{competition_path}/train.parquet"
            if os.path.exists(train_path):
                train_df = pd.read_parquet(train_path)
            else:
                train_df = pd.read_csv(f"{competition_path}/train.csv")
            test_path = f"{competition_path}/test.parquet"
            if os.path.exists(test_path):
                test_df = pd.read_parquet(test_path)
            else:
                test_df = pd.read_csv(f"{competition_path}/test.csv")
            print(f"Training: {train_df.shape}, Test: {test_df.shape}")
            return train_df, test_df
        except Exception as e:
            print("Data load error:", e)
            # Fallback small synthetic
            n, m = 1500, 25
            X = np.random.randn(n, m)
            y = X[:, :6].sum(axis=1) + 0.1*np.random.randn(n)
            train_df = pd.DataFrame(X, columns=[f"f{i}" for i in range(m)])
            train_df['target'] = y
            Xt = np.random.randn(300, m)
            test_df = pd.DataFrame(Xt, columns=[f"f{i}" for i in range(m)])
            test_df['id'] = range(len(test_df))
            return train_df, test_df

    def feature_engineering(self, df):
        df = df.copy()
        num_cols = [c for c in df.columns if c not in ['target','id'] and pd.api.types.is_numeric_dtype(df[c])]
        if num_cols:
            df['mean_all'] = df[num_cols].mean(axis=1)
            df['std_all'] = df[num_cols].std(axis=1)
            df['max_all'] = df[num_cols].max(axis=1)
            df['min_all'] = df[num_cols].min(axis=1)
        return df

class BoostingEnsemble:
    def __init__(self):
        self.models = {}
        self.scaler = None

    def train(self, X, y):
        X_tr, X_val, y_tr, y_val = train_test_split(X, y, test_size=0.2, random_state=42)
        self.scaler = StandardScaler(with_mean=True, with_std=True)
        X_trs = self.scaler.fit_transform(X_tr)
        X_vals = self.scaler.transform(X_val)

        # XGBoost
        print("Training XGBoost...")
        xgb_model = xgb.XGBRegressor(
            n_estimators=800, max_depth=8, learning_rate=0.05,
            subsample=0.85, colsample_bytree=0.8, random_state=42,
            tree_method='hist'
        )
        xgb_model.fit(X_trs, y_tr, eval_set=[(X_vals, y_val)], verbose=False)
        xgb_rmse = mean_squared_error(y_val, xgb_model.predict(X_vals), squared=False)
        print(f"XGB RMSE: {xgb_rmse:.6f}")
        self.models['xgb'] = xgb_model

        # LightGBM
        print("Training LightGBM...")
        lgb_model = lgb.LGBMRegressor(
            n_estimators=800, max_depth=-1, learning_rate=0.05,
            subsample=0.85, colsample_bytree=0.8, random_state=42,
        )
        lgb_model.fit(X_trs, y_tr, eval_set=[(X_vals, y_val)], verbose=False)
        lgb_rmse = mean_squared_error(y_val, lgb_model.predict(X_vals), squared=False)
        print(f"LGB RMSE: {lgb_rmse:.6f}")
        self.models['lgb'] = lgb_model

        # Gradient Boosting (baseline)
        print("Training Sklearn GB...")
        gb_model = GradientBoostingRegressor(n_estimators=500, max_depth=8, learning_rate=0.05, random_state=42)
        gb_model.fit(X_trs, y_tr)
        gb_rmse = mean_squared_error(y_val, gb_model.predict(X_vals), squared=False)
        print(f"GB RMSE: {gb_rmse:.6f}")
        self.models['gb'] = gb_model

    def predict(self, X):
        Xs = self.scaler.transform(X)
        preds = {
            name: model.predict(Xs) for name, model in self.models.items()
        }
        # Weighted average
        w = {'xgb':0.4, 'lgb':0.4, 'gb':0.2}
        final = np.zeros(len(X))
        for name, p in preds.items():
            final += w[name]*p
        return final

def main():
    print("="*80)
    print("ARIEL SPACE CHALLENGE 2025 - AUTOMATED EXECUTION (v3 XGBoost+LightGBM)")
    print("="*80)

    optimizer = CloudMLEOptimizer()
    proc = AdvancedSpaceDataProcessor()

    print("\n" + "="*60)
    print("DATA LOADING PHASE")
    train_df, test_df = proc.load_data_efficiently()

    print("\n" + "="*60)
    print("FEATURE ENGINEERING PHASE")
    train_df = proc.feature_engineering(train_df)
    test_df = proc.feature_engineering(test_df)

    target = 'target'
    id_col = 'id'
    feat_cols = [c for c in train_df.columns if c not in [target, id_col]]

    X = train_df[feat_cols].values
    y = train_df[target].values if target in train_df.columns else train_df.iloc[:,-1].values
    Xt = test_df[feat_cols].values

    print("\n" + "="*60)
    print("MODEL TRAINING PHASE")
    ens = BoostingEnsemble()
    ens.train(X, y)

    print("\n" + "="*60)
    print("PREDICTION PHASE")
    preds = ens.predict(Xt)

    print("\n" + "="*60)
    print("SUBMISSION CREATION")
    if id_col in test_df.columns:
        sub = pd.DataFrame({'id': test_df[id_col], 'target': preds})
    else:
        sub = pd.DataFrame({'id': np.arange(len(preds)), 'target': preds})
    sub.to_csv('submission.csv', index=False)

    print(f"Submission created: {len(sub)} rows")
    print("Done.")

    return sub

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