import os
import time
import math
import pickle
import torch
import torch.nn as nn
import pandas as pd
import numpy as np
from contextlib import asynccontextmanager
from typing import List, Optional
from datetime import datetime
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

# --- Schema Definitions ---

class WeatherPoint(BaseModel):
    timestamp: datetime = Field(..., description="Timestamp of the observation")
    temperature: float = Field(..., description="Temperature in Celsius (temperature_2m)")
    humidity: float = Field(..., description="Relative Humidity in % (relative_humidity_2m)")
    wind_speed: float = Field(..., description="Wind Speed in km/h (windspeed_10m)")
    
    class Config:
        json_schema_extra = {
            "example": {
                "timestamp": "2024-01-01T10:00:00",
                "temperature": 25.4,
                "humidity": 45.0,
                "wind_speed": 12.5
            }
        }

class PredictionRequest(BaseModel):
    features: List[WeatherPoint] = Field(..., min_items=24, max_items=24, description="List of exactly 24 weather points (last 24 hours)")
    historical_loads: Optional[List[float]] = Field(None, description="Historical load data (Ignored by Digital Twin models)")

class PredictionResponse(BaseModel):
    load: float = Field(..., description="Predicted System Load in MW")
    confidence_interval: Optional[List[float]] = Field(None, description="[Lower, Upper] bound of prediction confidence")
    model_name: str = Field(..., description="Name of the model used")
    execution_time: float = Field(..., description="Inference time in seconds")


# --- Model Architecture Definitions ---

class StandardLSTM(nn.Module):
    def __init__(self, input_size, hidden_size=64, num_layers=1, output_size=1, dropout=0.5):
        super(StandardLSTM, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=0 if num_layers==1 else dropout)
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        out, _ = self.lstm(x)
        out = self.dropout(out[:, -1, :])
        out = self.fc(out)
        return out

class DeepLSTM(nn.Module):
    def __init__(self, input_size, hidden_size=64, num_layers=2, output_size=1, dropout=0.2):
        super(DeepLSTM, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h0 = torch.zeros(self.lstm.num_layers, x.size(0), self.lstm.hidden_size).to(x.device)
        c0 = torch.zeros(self.lstm.num_layers, x.size(0), self.lstm.hidden_size).to(x.device)
        out, _ = self.lstm(x, (h0,c0)) 
        out = self.fc(out[:, -1, :])
        return out

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

class IndiaTransformer(nn.Module):
    def __init__(self, input_size, d_model, nhead, num_encoder_layers, dim_feedforward, output_size, dropout=0.1):
        super(IndiaTransformer, self).__init__()
        self.embedding = nn.Linear(input_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model, dropout)
        encoder_layers = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, 
                                                  dim_feedforward=dim_feedforward, 
                                                  dropout=dropout, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_encoder_layers)
        self.fc = nn.Linear(d_model, output_size)
        self.d_model = d_model

    def forward(self, src):
        src = self.embedding(src) * math.sqrt(self.d_model)
        src = self.pos_encoder(src)
        output = self.transformer_encoder(src)
        output = self.fc(output[:, -1, :])
        return output

# --- Preprocessor ---

A = 17.27
B = 237.7

class Preprocessor:
    def __init__(self, scaler_path_lstm: str, scaler_path_transformer: str):
        self.scaler_lstm = self._load_scaler(scaler_path_lstm)
        self.scaler_transformer = self._load_scaler(scaler_path_transformer)
        self.FEATURES = ["hour", "day_of_month", "dayofweek", "day_of_year", "month", "year", 
                         "week_of_year", "temperature_2m", "relative_humidity_2m", 
                         "windspeed_10m", "dew_point_2m"]

    def _load_scaler(self, path: str):
        if not os.path.exists(path):
            raise FileNotFoundError(f"Scaler not found at {path}")
        with open(path, 'rb') as f:
            return pickle.load(f)

    def calculate_dew_point(self, temp, humidity):
        alpha = ((A * temp) / (B + temp)) + math.log(humidity / 100.0)
        return (B * alpha) / (A - alpha)

    def prepare_input(self, weather_points: List[WeatherPoint], model_type: str = "lstm"):
        data = []
        for wp in weather_points:
            data.append({
                "timestamp": wp.timestamp,
                "temperature_2m": wp.temperature,
                "relative_humidity_2m": wp.humidity,
                "windspeed_10m": wp.wind_speed
            })
        df = pd.DataFrame(data)

        if df['timestamp'].dt.tz is not None:
            df['timestamp'] = df['timestamp'].dt.tz_convert('Asia/Kolkata')
            
        df['hour'] = df['timestamp'].dt.hour
        df['day_of_month'] = df['timestamp'].dt.day
        df['dayofweek'] = df['timestamp'].dt.dayofweek
        df['day_of_year'] = df['timestamp'].dt.dayofyear
        df['month'] = df['timestamp'].dt.month
        df['year'] = df['timestamp'].dt.year
        df['week_of_year'] = df['timestamp'].dt.isocalendar().week.astype(int)

        df['dew_point_2m'] = df.apply(lambda x: self.calculate_dew_point(x['temperature_2m'], x['relative_humidity_2m']), axis=1)
        df_features = df[self.FEATURES]
        
        vals = df_features.values
        vals_padded = np.hstack([vals, np.zeros((vals.shape[0], 1))])
        scaler = self.scaler_transformer if model_type == "transformer" else self.scaler_lstm
        vals_scaled = scaler.transform(vals_padded)
        
        X = vals_scaled[:, :-1]
        X = X.reshape(1, len(weather_points), 11)
        
        return torch.FloatTensor(X)


# --- Global State ---

models = {}
preprocessor = None
device = torch.device('cpu')

@asynccontextmanager
async def lifespan(app: FastAPI):
    global models, preprocessor
    base_dir = os.path.dirname(__file__)
    
    # Paths
    std_lstm_path = os.path.join(base_dir, "india_models", "standard_lstm.pt")
    deep_lstm_path = os.path.join(base_dir, "india_models", "deep_lstm.pt")
    scaler_lstm_path = os.path.join(base_dir, "india_models", "scaler.pkl")
    
    transformer_path = os.path.join(base_dir, "transformer_model", "best_transformer.pt")
    scaler_transformer_path = os.path.join(base_dir, "transformer_model", "scaler_transformer.pkl")
    
    print("Loading preprocessor...")
    preprocessor = Preprocessor(scaler_lstm_path, scaler_transformer_path)
    
    print("Loading Standard LSTM...")
    std_lstm = StandardLSTM(input_size=11, hidden_size=64, num_layers=1, output_size=1)
    std_lstm.load_state_dict(torch.load(std_lstm_path, map_location=device, weights_only=True))
    std_lstm.to(device)
    std_lstm.eval()
    models['lstm'] = std_lstm

    print("Loading Deep LSTM...")
    deep_lstm = DeepLSTM(input_size=11, hidden_size=64, num_layers=2, output_size=1)
    deep_lstm.load_state_dict(torch.load(deep_lstm_path, map_location=device, weights_only=True))
    deep_lstm.to(device)
    deep_lstm.eval()
    models['deeplstm'] = deep_lstm

    print("Loading Transformer...")
    transformer = IndiaTransformer(
        input_size=11, d_model=64, nhead=4, 
        num_encoder_layers=2, dim_feedforward=128, output_size=1
    )
    transformer.load_state_dict(torch.load(transformer_path, map_location=device, weights_only=True))
    transformer.to(device)
    transformer.eval()
    models['transformer'] = transformer

    print("Startup complete.")
    yield
    models.clear()
    print("Shutdown complete.")

app = FastAPI(title="GridSim Inference API", lifespan=lifespan)

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

def infer(request: PredictionRequest, model_key: str, model_type: str) -> PredictionResponse:
    start_time = time.time()
    
    if len(request.features) != 24:
        raise HTTPException(status_code=400, detail="Exactly 24 features required")
        
    if model_key not in models:
        raise HTTPException(status_code=500, detail=f"Model {model_key} not loaded")
        
    model = models[model_key]
    
    try:
        input_tensor = preprocessor.prepare_input(request.features, model_type).to(device)
        
        with torch.no_grad():
            output_scaled = model(input_tensor)
            
        scaler = preprocessor.scaler_transformer if model_type == "transformer" else preprocessor.scaler_lstm
        
        # Unscale
        # Output is just 1 value, pad with 11 zeros for features
        dummy = np.zeros((1, 12))
        dummy[0, -1] = output_scaled.item()
        unscaled_pred = scaler.inverse_transform(dummy)[0, -1]
        
        exec_time = time.time() - start_time
        
        # Simple dummy CI for now as not specified
        ci_lower = unscaled_pred * 0.95
        ci_upper = unscaled_pred * 1.05
        
        return PredictionResponse(
            load=float(unscaled_pred),
            confidence_interval=[float(ci_lower), float(ci_upper)],
            model_name=model_key,
            execution_time=exec_time
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/predict/lstm", response_model=PredictionResponse)
def predict_lstm(request: PredictionRequest):
    return infer(request, 'lstm', 'lstm')

@app.post("/predict/deeplstm", response_model=PredictionResponse)
def predict_deeplstm(request: PredictionRequest):
    return infer(request, 'deeplstm', 'lstm')

@app.post("/predict/transformer", response_model=PredictionResponse)
def predict_transformer(request: PredictionRequest):
    return infer(request, 'transformer', 'transformer')

@app.get("/health")
def health():
    return {"status": "ok", "models_loaded": list(models.keys())}

@app.get("/model/info")
def model_info():
    return {
        "models": {
            "lstm": "Standard LSTM (1 layer, hidden=64)",
            "deeplstm": "Deep LSTM (2 layers, hidden=64)",
            "transformer": "Transformer (3 layers, d_model=64, nhead=4)"
        },
        "features_expected": 24,
        "device": str(device)
    }