import pandas as pd
import numpy as np
import os
from glob import glob
from tqdm import tqdm

#%%
def step_one()
 print("Step 1")
"""
Grouping Data files by Instruments 
"""

base_dir = "/kaggle/input/ariel-data-challenge-2025/train"
planet_ids = sorted(os.listdir(base_dir), key=lambda x: int(x))

def extract_signal_data(instrument):
    adc_df = pd.read_csv("/kaggle/input/ariel-data-challenge-2025/adc_info.csv")
    gain = adc_df[f"{instrument}_adc_gain"].iloc[0]
    offset = adc_df[f"{instrument}_adc_offset"].iloc[0]

    features = []
    used_ids = []

    print(f"Processing {instrument} signals...")
    for pid in tqdm(planet_ids):
        signal_files = glob(f"{base_dir}/{pid}/{instrument}_signal_*.parquet")
        mean_rows = []

        for path in signal_files:
            try:
                arr = pd.read_parquet(path).values.astype('float32')
                arr = arr * gain + offset
                mean_rows.append(np.mean(arr, axis=0))
            except:
                continue

        if mean_rows:
            features.append(np.mean(mean_rows, axis=0))
            used_ids.append(int(pid))

    df = pd.DataFrame(features, index=used_ids)
    df.to_parquet(f"{instrument.lower()}_signal_all.parquet")
    return df

# Extract and save both signal types
airs_signal = extract_signal_data("AIRS-CH0")
fgs_signal = extract_signal_data("FGS1")

print("✅ AIRS-CH0 and FGS1 signal files saved.")

# === General function to extract and save calibration frames ===
def extract_calibration_data(instrument, calib_type):
    all_frames = []
    used_planets = []

    print(f"Processing {instrument} - {calib_type}...")
    for pid in tqdm(planet_ids):
        path = os.path.join(base_dir, pid, f"{instrument}_calibration_0", f"{calib_type}.parquet")
        if os.path.exists(path):
            try:
                df = pd.read_parquet(path).values.flatten()
                all_frames.append(df)
                used_planets.append(int(pid))
            except:
                continue

    df = pd.DataFrame(all_frames, index=used_planets)
    df.to_parquet(f"{instrument.lower()}_{calib_type}_all.parquet")
    return df

# === Extract and save AIRS-CH0 calibration files ===
airs_dark   = extract_calibration_data("AIRS-CH0", "dark")
airs_flat   = extract_calibration_data("AIRS-CH0", "flat")
airs_dead   = extract_calibration_data("AIRS-CH0", "dead")
airs_linear = extract_calibration_data("AIRS-CH0", "linear_corr")
airs_read   = extract_calibration_data("AIRS-CH0", "read")

# === Extract and save FGS1 calibration files ===
fgs_dark   = extract_calibration_data("FGS1", "dark")
fgs_flat   = extract_calibration_data("FGS1", "flat")
fgs_dead   = extract_calibration_data("FGS1", "dead")
fgs_linear = extract_calibration_data("FGS1", "linear_corr")
fgs_read   = extract_calibration_data("FGS1", "read")

print("✅ AIRS-CH0 and FGS1 calibration files saved.")

#%%
def step_two()
 print("Step 2")

# === PHASE 1: Model Training and Evaluation ===

from sklearn.linear_model import Ridge
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
from sklearn.preprocessing import StandardScaler

# === Load input features and labels ===
X = pd.read_parquet("/kaggle/working/airs-ch0_signal_all.parquet")  # Already stacked AIRS-CH0 train features
y_df = pd.read_csv("/kaggle/input/ariel-data-challenge-2025/train.csv")  # Labels: 283 spectrum values

# Ensure perfect planet_id alignment between X and y
X = X.sort_index()
y = y_df.set_index("planet_id").loc[X.index]

# === Optional: Standardize features ===
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# === Train/validation split ===
X_train, X_val, y_train, y_val = train_test_split(X_scaled, y, test_size=0.15, random_state=42)

# === Train Ridge Regression ===
ridge = Ridge(alpha=1.0)
ridge.fit(X_train, y_train)

# === Predict on validation set ===
y_pred = ridge.predict(X_val)

# === Evaluation ===
rmse = np.sqrt(mean_squared_error(y_val, y_pred))
print(f"Validation RMSE: {rmse:.5f}")

# === PHASE 2: Generate Predictions for Test Planets ===

# Load test AIRS-CH0 signal features (must match same shape/processing as training)
X_test = pd.read_parquet("/kaggle/working/airs-ch0_signal_test_all.parquet")  # Your preprocessed test file

# Ensure test is sorted to match planet_ids correctly
X_test = X_test.sort_index()

# Apply the same scaler used on training set
X_test_scaled = scaler.transform(X_test)

# Predict mean spectrum μ for test set
mu_pred = ridge.predict(X_test_scaled)  # Shape: (n_test, 283)

# === Estimate uncertainty σ for each wavelength ===
# Simple estimate: Use std of residuals on train set
train_residuals = y_train.values - ridge.predict(X_train)
residual_std = np.std(train_residuals, axis=0)  # Shape: (283,)
sigma_pred = np.tile(residual_std, (mu_pred.shape[0], 1))  # Repeat std per row

# Optional: Add a small value to avoid zero uncertainty
sigma_pred += 1e-6

# Check shape
print("μ shape:", mu_pred.shape)
print("σ shape:", sigma_pred.shape)

# === PHASE 3: Create submission.csv in proper format ===

# Ensure planet IDs are in order matching test set
planet_ids = X_test.index.tolist()

# Create column names for μ (mean prediction) and σ (uncertainty)
mu_cols = [f"mu_{i}" for i in range(283)]
sigma_cols = [f"sigma_{i}" for i in range(283)]

# Build the submission DataFrame
submission = pd.DataFrame(mu_pred, columns=mu_cols)
submission["planet_id"] = planet_ids  # First column

# Add uncertainty (σ) predictions
for i, col in enumerate(sigma_cols):
    submission[col] = sigma_pred[:, i]

# Reorder columns: planet_id, mu_0...mu_282, sigma_0...sigma_282
submission = submission[["planet_id"] + mu_cols + sigma_cols]

# Save to CSV
submission.to_csv("submission.csv", index=False)

print("✅ submission.csv created successfully!")
print("Shape:", submission.shape)
print("Preview:")
print(submission.head(2))