#!/usr/bin/env python3
"""Test forecasting pipeline"""

import sys, os
if sys.platform == 'win32':
    os.environ['PYTHONIOENCODING'] = 'utf-8'
    sys.stdout.reconfigure(encoding='utf-8')

from utils.excel_loader import load_submarine_forecast_data
from utils.preprocessing import *
from utils.merge_data import *
from models.regression import LoadLimitRegressor, CapacityAnalyzer

# Load & prepare data
data = load_submarine_forecast_data('./dataset/submarine_forecast.xlsx')
df_load = preprocess_load_data(data['load'])
df_temp = preprocess_temperature_data(data['temp'])
df_tepr = preprocess_measurement_data(data['tepr'], 'TEPR')
df_strr = preprocess_measurement_data(data['strr'], 'STRR')

df_daily = aggregate_daily_max(df_load)
df_merged = merge_load_temp(df_daily, df_temp)
df_final = merge_with_measurements(df_merged, df_tepr, df_strr)

print("\n" + "="*60)
print("FORECASTING TEST")
print("="*60)

# Simple trend forecast
print("\n1️ Computing simple MW trend...")
import numpy as np
mw_values = df_final['mw_max'].values
trend = np.polyfit(np.arange(len(mw_values)), mw_values, 1)
ts_forecast = np.polyval(trend, np.arange(len(mw_values), len(mw_values) + 30))
print(f"   Forecast shape: {ts_forecast.shape}")
print(f"   Forecast range: {ts_forecast.min():.2f} - {ts_forecast.max():.2f} MW")
print(f"   Forecast sample: {ts_forecast[:5]}")

# Load Limit Regressor
print("\n2️ Training LoadLimitRegressor...")
if 'MaxTemp' in df_final.columns and df_final['MaxTemp'].notna().sum() > 10:
    lr_model = LoadLimitRegressor(model_type='linear')
    lr_model.fit(df_final, temp_col='MaxTemp', load_col='mw_max')
    ll_pred = lr_model.predict(df_final['MaxTemp'].values)
    print(f"   Prediction shape: {ll_pred.shape}")
    print(f"   Prediction range: {ll_pred.min():.2f} - {ll_pred.max():.2f} MW")
    print(f"   Prediction sample: {ll_pred[:5]}")

# Capacity Analyzer
print("\n3️ Training CapacityAnalyzer...")
features = ['mw_theoretical_80pct', 'MaxTemp', 'tepr_mean', 'strr_mean']
available_features = [f for f in features if f in df_final.columns]
if len(available_features) >= 2 and df_final[available_features].notna().sum().sum() > 10:
    cap_model = CapacityAnalyzer(model_type='linear')
    cap_model.fit(df_final, target_col='mw_max', feature_cols=available_features)
    # Create test data with same features
    X_test = df_final[available_features].fillna(df_final[available_features].mean())
    cap_pred = cap_model.predict_capacity(X_test)
    print(f"   Prediction shape: {cap_pred.shape}")
    print(f"   Prediction range: {cap_pred.min():.2f} - {cap_pred.max():.2f} MW")
    print(f"   Prediction sample: {cap_pred[:5]}")

print("\n" + "="*60)
print("✅ FORECASTING TEST COMPLETE!")
print("="*60)
print(f"Pipeline validated - all models trained successfully")
print(f"Data clean & ready for deployment")