# ============================================================
# 📈 app.py — Backend FastAPI para predicciones BBVA y SANTANDER
# ============================================================

from fastapi import FastAPI
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
import torch
from torch import nn
import joblib
import numpy as np
import pandas as pd

# ============================================================
# 1️⃣ Configuración base de FastAPI
# ============================================================
app = FastAPI(title="Predicciones BBVA y SANTANDER API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ============================================================
# 2️⃣ Definir ambos modelos (LSTM y GRU)
# ============================================================
class LSTMModel(nn.Module):
    def __init__(self, n_features, hidden_size=64):
        super().__init__()
        self.lstm = nn.LSTM(input_size=n_features, hidden_size=hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, 1)
    def forward(self, x):
        out, _ = self.lstm(x)
        out = out[:, -1, :]
        out = self.fc(out)
        return out.squeeze()


class GRUModel(nn.Module):
    def __init__(self, n_features, hidden_size=64, num_layers=1):
        super().__init__()
        self.gru = nn.GRU(
            input_size=n_features,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True
        )
        self.fc = nn.Linear(hidden_size, 1)
    def forward(self, x):
        out, _ = self.gru(x)
        out = out[:, -1, :]
        out = self.fc(out)
        return out.squeeze()

# ============================================================
# 3️⃣ Cargar ambos modelos y scalers
# ============================================================
device = torch.device("cpu")

# --- BBVA ---
bbva_model = LSTMModel(n_features=5, hidden_size=64)
bbva_model.load_state_dict(torch.load("models_bbva/bbva_lstm_model.pth", map_location=device))
bbva_model.to(device)
bbva_model.eval()
bbva_scaler_features = joblib.load("models_bbva/scaler_features_full.pkl")
bbva_scaler_target = joblib.load("models_bbva/scaler_target.pkl")

# --- SANTANDER ---
santander_model = GRUModel(n_features=5, hidden_size=64)
santander_model.load_state_dict(torch.load("models_santander/santander_gru_model.pth", map_location=device))
santander_model.to(device)
santander_model.eval()
santander_scaler_features = joblib.load("models_santander/santander_scaler_features_full.pkl")
santander_scaler_target = joblib.load("models_santander/santander_scaler_target.pkl")

# ============================================================
# 4️⃣ Modelo de entrada
# ============================================================
class PredictRequest(BaseModel):
    n_future: int = 3

# ============================================================
# 5️⃣ Función genérica para predicción
# ============================================================
def generar_prediccion(df, model, scaler_features, scaler_target, features, cols_to_scale, window_size, n_future):
    # Escalar
    df_scaled = df.copy()
    df_scaled[cols_to_scale] = scaler_features.transform(df[cols_to_scale])

    # Última ventana
    last_window = df_scaled[features].iloc[-window_size:].values

    preds_scaled = []
    window = last_window.copy()
    for _ in range(n_future):
        X_input = torch.tensor(window, dtype=torch.float32).unsqueeze(0)
        with torch.no_grad():
            pred_scaled = model(X_input).cpu().numpy().flatten()[0]
        preds_scaled.append(pred_scaled)
        next_day = window[-1].copy()
        window = np.vstack([window[1:], next_day])

    # Desescalar
    preds_real = scaler_target.inverse_transform(np.array(preds_scaled).reshape(-1, 1)).flatten().tolist()

    # Fechas futuras
    last_date = df.index[-1]
    future_dates = pd.date_range(last_date + pd.Timedelta(days=1), periods=n_future, freq="B")
    fechas = [str(d.date()) for d in future_dates]

    return fechas, preds_real

# ============================================================
# 6️⃣ Endpoint raíz
# ============================================================
@app.get("/")
def home():
    return {"message": "✅ API de predicciones BBVA y Santander operativa"}

# ============================================================
# 7️⃣ Endpoint principal
# ============================================================
@app.post("/predict/{banco}")
def predict(banco: str, req: PredictRequest):
    banco = banco.lower()
    n_future = req.n_future

    # === BBVA ===
    if banco == "bbva":
        path = "data/bbva_pred.csv"
        df = pd.read_csv(path)
        df["Date"] = pd.to_datetime(df["Date"])
        df = df.set_index("Date")
        df = df.loc["2021-01-01":"2025-10-31"]

        cols_to_scale = [
            'Open','High','Low','Close','Adj Close',
            'ma_5','close_lag1','close_lag2','close_lag3',
            'Volume','ibex_momentum_5d'
        ]
        features = ['Open','High','Low','return_1d','return_3d']

        fechas, preds = generar_prediccion(df, bbva_model, bbva_scaler_features, bbva_scaler_target, features, cols_to_scale, 3, n_future)
        return {"banco": "BBVA", "fechas": fechas, "predicciones": preds}

    # === SANTANDER ===
    elif banco == "santander":
        path = "data/santander_pred.csv"
        df = pd.read_csv(path)
        df["Date"] = pd.to_datetime(df["Date"])
        df = df.set_index("Date")
        df = df.loc["2021-01-01":"2025-10-31"]

        cols_to_scale = [
            'Open','High','Low','Close','Adj Close',
            'ma_5','close_lag1','close_lag2','close_lag3',
            'Volume','ibex_momentum_5d'
        ]
        features = ['Open','High','Low','close_lag1','close_lag2']

        fechas, preds = generar_prediccion(df, santander_model, santander_scaler_features, santander_scaler_target, features, cols_to_scale, 3, n_future)
        return {"banco": "Santander", "fechas": fechas, "predicciones": preds}

    # === Otro banco no válido ===
    else:
        return {"error": "Banco no reconocido. Usa 'bbva' o 'santander'."}

# ============================================================
# 8️⃣ Nuevo endpoint de simulación — Variar "Open"
# ============================================================

class SimulateRequest(BaseModel):
    factor_open: float = 1.0  # por defecto, sin cambio
    n_future: int = 3         # número de días de predicción


@app.post("/simulate/{banco}")
def simulate(banco: str, req: SimulateRequest):
    banco = banco.lower()
    n_future = req.n_future
    factor_open = req.factor_open

    # === BBVA ===
    if banco == "bbva":
        path = "data/bbva_pred.csv"
        df = pd.read_csv(path)
        df["Date"] = pd.to_datetime(df["Date"])
        df = df.set_index("Date")
        df = df.loc["2021-01-01":"2025-10-31"]

        cols_to_scale = [
            'Open','High','Low','Close','Adj Close',
            'ma_5','close_lag1','close_lag2','close_lag3',
            'Volume','ibex_momentum_5d'
        ]
        features = ['Open','High','Low','return_1d','return_3d']

        # Escalar los datos
        df_scaled = df.copy()
        df_scaled[cols_to_scale] = bbva_scaler_features.transform(df[cols_to_scale])

        # Tomar la última ventana y modificar solo "Open"
        last_window = df_scaled[features].iloc[-3:].values
        modified_window = last_window.copy()
        modified_window[:, 0] *= factor_open  # multiplicar solo el "Open"

        preds_scaled = []
        window = modified_window.copy()
        for _ in range(n_future):
            X_input = torch.tensor(window, dtype=torch.float32).unsqueeze(0)
            with torch.no_grad():
                pred_scaled = bbva_model(X_input).cpu().numpy().flatten()[0]
            preds_scaled.append(pred_scaled)
            next_day = window[-1].copy()
            window = np.vstack([window[1:], next_day])

        preds_real = bbva_scaler_target.inverse_transform(np.array(preds_scaled).reshape(-1, 1)).flatten().tolist()

        last_date = df.index[-1]
        future_dates = pd.date_range(last_date + pd.Timedelta(days=1), periods=n_future, freq="B")
        fechas = [str(d.date()) for d in future_dates]

        return {
            "banco": "BBVA",
            "factor_open": factor_open,
            "fechas": fechas,
            "predicciones": preds_real
        }

    # === SANTANDER ===
    elif banco == "santander":
        path = "data/santander_pred.csv"
        df = pd.read_csv(path)
        df["Date"] = pd.to_datetime(df["Date"])
        df = df.set_index("Date")
        df = df.loc["2021-01-01":"2025-10-31"]

        cols_to_scale = [
            'Open','High','Low','Close','Adj Close',
            'ma_5','close_lag1','close_lag2','close_lag3',
            'Volume','ibex_momentum_5d'
        ]
        features = ['Open','High','Low','close_lag1','close_lag2']

        df_scaled = df.copy()
        df_scaled[cols_to_scale] = santander_scaler_features.transform(df[cols_to_scale])

        last_window = df_scaled[features].iloc[-3:].values
        modified_window = last_window.copy()
        modified_window[:, 0] *= factor_open

        preds_scaled = []
        window = modified_window.copy()
        for _ in range(n_future):
            X_input = torch.tensor(window, dtype=torch.float32).unsqueeze(0)
            with torch.no_grad():
                pred_scaled = santander_model(X_input).cpu().numpy().flatten()[0]
            preds_scaled.append(pred_scaled)
            next_day = window[-1].copy()
            window = np.vstack([window[1:], next_day])

        preds_real = santander_scaler_target.inverse_transform(np.array(preds_scaled).reshape(-1, 1)).flatten().tolist()

        last_date = df.index[-1]
        future_dates = pd.date_range(last_date + pd.Timedelta(days=1), periods=n_future, freq="B")
        fechas = [str(d.date()) for d in future_dates]

        return {
            "banco": "Santander",
            "factor_open": factor_open,
            "fechas": fechas,
            "predicciones": preds_real
        }

    # === Error ===
    else:
        return {"error": "Banco no reconocido. Usa 'bbva' o 'santander'."}