import yfinance as yf
import pandas as pd
import numpy as np
from scipy.optimize import differential_evolution
import json
import os
import ta

# Import your existing logic
from nifty100 import NIFTY100
from backtest_engine import backtest_strategy
from ml_model import predict_probability
from monte_carlo import monte_carlo_probability
from technical_agent import analyze_technical

WEIGHTS_FILE = "ticker_weights.json"

def simulate_historical_features(df, lookback_days=30):
    """
    Pre-computes features and forward returns for optimization to save compute.
    Instead of 4 years of daily compute (which would crash), we sample recent history.
    """
    features = []
    
    # Ensure we have enough data to slice
    if len(df) < 150:
        return features

    # Test the last `lookback_days` to find optimal weights
    start_idx = len(df) - lookback_days
    
    for i in range(start_idx, len(df) - 5): # Leave 5 days for forward return checking
        
        sliced_df = df.iloc[:i].copy()
        current_price = float(sliced_df["Close"].iloc[-1])
        
        # Calculate ATR for targets/stops
        high = sliced_df["High"].squeeze()
        low = sliced_df["Low"].squeeze()
        close = sliced_df["Close"].squeeze()
        
        atr = ta.volatility.AverageTrueRange(high=high, low=low, close=close, window=14).average_true_range().iloc[-1]
        
        stop = round(current_price - (1.5 * atr), 2)
        target = round(current_price + (2.5 * atr), 2) # simplified mid-tier target
        risk_reward = round((target - current_price) / (current_price - stop), 2) if (current_price - stop) != 0 else 0
        
        # Generate model probabilities
        ml_prob = predict_probability(sliced_df)
        mc_prob = monte_carlo_probability(sliced_df, target, stop, simulations=200, horizon=5) # Reduced for speed
        winrate = backtest_strategy(sliced_df)
        
        # Approximate base confidence (Technicals)
        dummy_dict = {"TICKER": sliced_df}
        tech_score = analyze_technical(dummy_dict).get("TICKER", 0)
        confidence = max(0, min(100, round((tech_score / 5) * 100, 2))) # simplified
        
        # Check if trade was actually a win in the future
        future_df = df.iloc[i:i+5]
        hit_target = any(future_df["High"] >= target)
        hit_stop = any(future_df["Low"] <= stop)
        
        is_win = 1 if hit_target and not hit_stop else 0
        
        features.append({
            "confidence": confidence,
            "winrate": winrate,
            "risk_reward": risk_reward,
            "ml_prob": ml_prob,
            "mc_prob": mc_prob,
            "is_win": is_win
        })
        
    return features

def objective_function(weights, features):
    """
    Calculates the negative win rate (since scipy minimizes) for a given set of weights.
    Trigger threshold is set to 60.
    """
    w1, w2, w3, w4, w5 = weights
    
    trades_triggered = 0
    winning_trades = 0
    
    for f in features:
        score = (
            f["confidence"] * w1 +
            f["winrate"] * w2 +
            f["risk_reward"] * w3 +
            f["ml_prob"] * w4 +
            f["mc_prob"] * w5
        )
        
        if score > 60: # Trigger threshold
            trades_triggered += 1
            if f["is_win"]:
                winning_trades += 1
                
    if trades_triggered == 0:
        return 0 # Neutral penalty if no trades ever trigger
        
    win_rate = winning_trades / trades_triggered
    return -win_rate # Negative because we want to maximize

def optimize_all_tickers():
    
    best_weights = {}
    
    # Bounds for the weights: (min, max)
    bounds = [(0.0, 1.0), (0.0, 1.0), (0.0, 20.0), (0.0, 1.0), (0.0, 1.0)]
    
    # For testing on Spaces, we'll limit to a few tickers. 
    # Run locally to do all 100.
    tickers_to_run =  NIFTY100
    
    for ticker in tickers_to_run:
        print(f"Optimizing {ticker}...")
        try:
            df = yf.download(ticker, period="4y", interval="1d", progress=False)
            df.columns = df.columns.get_level_values(0)
            
            features = simulate_historical_features(df, lookback_days=30)
            
            if not features:
                continue
                
            result = differential_evolution(
                objective_function, 
                bounds, 
                args=(features,), 
                maxiter=10, 
                popsize=5
            )
            
            # Save optimized weights or fallback to defaults if optimization failed
            if result.success:
                best_weights[ticker] = {
                    "w1": float(result.x[0]),
                    "w2": float(result.x[1]),
                    "w3": float(result.x[2]),
                    "w4": float(result.x[3]),
                    "w5": float(result.x[4])
                }
            
        except Exception as e:
            print(f"Failed to optimize {ticker}: {e}")
            pass
            
    with open(WEIGHTS_FILE, "w") as f:
        json.dump(best_weights, f, indent=4)
        
    # Change the return statement to pass the file path back
    return "Optimization Complete! Click below to download.", WEIGHTS_FILE

if __name__ == "__main__":
    optimize_all_tickers()