import os, json, math, pickle
from datetime import datetime, timedelta
import numpy as np
import yfinance as yf
from sklearn.metrics import mean_absolute_error, mean_squared_error
import torch

from models import StockLSTM

os.environ["CUDA_VISIBLE_DEVICES"] = ""
ARTIFACTS_DIR = "artifacts"

@torch.no_grad()
def evaluate(symbol: str):
    base = os.path.join(ARTIFACTS_DIR, symbol.upper())
    model = StockLSTM(input_dim=1, hidden_dim=64, num_layers=2, dropout=0.2)
    model.load_state_dict(torch.load(os.path.join(base, "model.pt"), map_location="cpu"))
    model.eval()
    with open(os.path.join(base, "scaler.pkl"), "rb") as f:
        scaler = pickle.load(f)
    with open(os.path.join(base, "meta.json"), "r") as f:
        meta = json.load(f)

    seq_len = meta["seq_len"]

    end = datetime.utcnow().date()
    start = end - timedelta(days=5*365)
    df = yf.download(symbol, start=start.isoformat(), end=end.isoformat(), progress=False, auto_adjust=True)
    data = df[["Close"]].dropna()
    
    # Compute returns
    data['LogReturn'] = np.log(data['Close'] / data['Close'].shift(1))
    data = data.dropna()
    
    returns = data['LogReturn'].values.reshape(-1, 1)

    scaled = scaler.transform(returns)
    split_idx = int(len(scaled) * 0.8)
    test_scaled = scaled[split_idx - seq_len:]  # include tail of train for continuity

    # build sequences
    X, y = [], []
    for i in range(seq_len, len(test_scaled)):
        X.append(test_scaled[i-seq_len:i])
        y.append(test_scaled[i])
    X = np.array(X, dtype=np.float32)
    y = np.array(y, dtype=np.float32)

    X_t = torch.from_numpy(X)       # [N, T, 1]
    pred_scaled = model(X_t).numpy()
    
    # inverse returns
    pred_returns = scaler.inverse_transform(pred_scaled).flatten()
    
    # Reconstruct prices
    # We need the price before the first test prediction
    # The test set in 'data' starts at split_idx
    # The first prediction corresponds to return at split_idx
    # So base is price at split_idx - 1
    
    # Note: 'data' here is the return-df (shifted). 
    # We need indices from the original df.
    # It's cleaner to just align by length.
    
    # Get original prices aligned with returns
    # df['Close'] has N+1 items if returns has N items.
    # data indices are a subset of df indices
    
    # Let's match by index
    test_indices = data.index[split_idx:]
    # Price predecessors (bases)
    # If a return is at time t, it depends on Price[t-1]
    
    # Simple reconstruction:
    # Get the price immediately preceding the test set
    base_price_idx = split_idx - 1
    if base_price_idx < 0:
        # Fallback if split is at 0 (unlikely)
        base_price = df['Close'].iloc[0]
    else:
        # The return at data.iloc[base_price_idx] is NOT the price
        # data only has returns.
        # We need to look at the original DF
        # The 'data' was created by dropping first row of df.
        # So data.iloc[0] corresponds to df.iloc[1].
        # data.iloc[split_idx] is roughly df.iloc[split_idx+1]
        
        # Exact alignment:
        # data index i matches df index i (if we kept index)
        pass

    # Let's rely on the original df
    # The returns in "y" (targets) correspond to `data.iloc[split_idx:]`
    # The Prices we want to compare against are `df['Close'][data.index[split_idx:]]`
    
    y_true_prices = df['Close'].loc[data.index[split_idx:]].values
    
    # Base price for the FIRST prediction:
    # The first return predicted is for data.index[split_idx]
    # So we need Price at data.index[split_idx-1] (previous day)
    # OR simpler: df['Close'].loc[data.index[split_idx-1]]
    
    first_test_idx_pos = df.index.get_loc(data.index[split_idx])
    base_price = df['Close'].iloc[first_test_idx_pos - 1]
    
    reconstructed = []
    curr = base_price
    for r in pred_returns:
        curr = curr * np.exp(r)
        reconstructed.append(curr)
        
    pred = np.array(reconstructed)
    y_true = y_true_prices[:len(pred)] # sync lengths

    rmse = math.sqrt(mean_squared_error(y_true, pred))
    mae = mean_absolute_error(y_true, pred)
    return {"symbol": symbol.upper(), "rmse": rmse, "mae": mae, "n": len(y_true)}