import os
import json
import pandas as pd
import numpy as np
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

DATA_FILE = os.path.join(os.path.dirname(__file__), "data", "nifty50_daily.parquet")
PREDICTIONS_FILE = os.path.join(os.path.dirname(__file__), "predictions.json")

def compute_rsi(series, window):
    delta = series.diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
    rs = gain / (loss + 1e-9)
    return 100 - (100 / (1 + rs))

def generate_predictions():
    if not os.path.exists(DATA_FILE):
        print(f"Data file missing: {DATA_FILE}")
        return
        
    df_all = pd.read_parquet(DATA_FILE)
    tickers = df_all['ticker'].unique()
    
    predictions = {}
    forecast_date = None
    
    for ticker in tickers:
        df = df_all[df_all['ticker'] == ticker].copy()
        df.sort_values('date', inplace=True)
        df.set_index('date', inplace=True)
        
        if len(df) < 150:
            continue
            
        forecast_date_ts = df.index[-1] + pd.Timedelta(days=1)
        if forecast_date_ts.weekday() >= 5:
            forecast_date_ts += pd.Timedelta(days=(7 - forecast_date_ts.weekday()))
            
        if forecast_date is None:
            forecast_date = forecast_date_ts.strftime('%Y-%m-%d')
            
        daily_close = df['close']
        df_feat = pd.DataFrame(index=df.index)
        df_feat['close'] = daily_close
        
        # Massive Feature Set
        for w in [2, 3, 5, 7, 14]:
            df_feat[f'rsi_{w}'] = compute_rsi(daily_close, w)
            
        for w in [3, 5, 10, 20]:
            df_feat[f'dist_sma_{w}'] = daily_close / daily_close.rolling(w).mean()
            
        for lag in [1, 2, 3, 5]:
            df_feat[f'ret_{lag}d'] = daily_close.pct_change(lag)
            
        df_feat.replace([np.inf, -np.inf], np.nan, inplace=True)
        df_feat = df_feat.dropna()
        
        if len(df_feat) < 130:
            continue
            
        # Target for historical testing
        df_feat['actual_dir'] = np.where(df_feat['close'].shift(-1) > df_feat['close'], 1, -1)
        
        # Test on 120 days BEFORE today to find the best rule
        test_set = df_feat.iloc[-121:-1]
        today_data = df_feat.iloc[-1]
        
        best_acc = 0
        best_rule = None
        
        # --- MASSIVE GRID SEARCH ---
        # 1. RSI Rules
        for w in [2, 3, 5, 7, 14]:
            feat = f'rsi_{w}'
            for thresh in range(10, 92, 2):
                sig_lt = np.where(test_set[feat] < thresh, 1, -1)
                acc_lt = (sig_lt == test_set['actual_dir']).mean()
                if acc_lt > best_acc:
                    best_acc = acc_lt
                    best_rule = (feat, thresh, '<')
                    
                sig_gt = np.where(test_set[feat] > thresh, 1, -1)
                acc_gt = (sig_gt == test_set['actual_dir']).mean()
                if acc_gt > best_acc:
                    best_acc = acc_gt
                    best_rule = (feat, thresh, '>')
                    
        # 2. SMA Distance Rules
        for w in [3, 5, 10, 20]:
            feat = f'dist_sma_{w}'
            for thresh in np.arange(0.85, 1.15, 0.005):
                sig_lt = np.where(test_set[feat] < thresh, 1, -1)
                acc_lt = (sig_lt == test_set['actual_dir']).mean()
                if acc_lt > best_acc:
                    best_acc = acc_lt
                    best_rule = (feat, thresh, '<')
                    
                sig_gt = np.where(test_set[feat] > thresh, 1, -1)
                acc_gt = (sig_gt == test_set['actual_dir']).mean()
                if acc_gt > best_acc:
                    best_acc = acc_gt
                    best_rule = (feat, thresh, '>')
                    
        # 3. Return Rules
        for lag in [1, 2, 3, 5]:
            feat = f'ret_{lag}d'
            for thresh in np.arange(-0.05, 0.052, 0.0025):
                sig_lt = np.where(test_set[feat] < thresh, 1, -1)
                acc_lt = (sig_lt == test_set['actual_dir']).mean()
                if acc_lt > best_acc:
                    best_acc = acc_lt
                    best_rule = (feat, thresh, '<')
                    
                sig_gt = np.where(test_set[feat] > thresh, 1, -1)
                acc_gt = (sig_gt == test_set['actual_dir']).mean()
                if acc_gt > best_acc:
                    best_acc = acc_gt
                    best_rule = (feat, thresh, '>')
                    
        # Apply the best rule found on test_set to TODAY
        feature, thresh, op = best_rule
        val = today_data[feature]
        
        if op == '<':
            prediction = 1 if val < thresh else -1
        else:
            prediction = 1 if val > thresh else -1
            
        # Format rule for display
        if 'dist_sma' in feature or 'ret_' in feature:
            rule_str = f"{feature} {op} {thresh:.4f}"
        else:
            rule_str = f"{feature} {op} {thresh}"
            
        predictions[ticker] = {
            "prediction": "UP" if prediction == 1 else "DOWN",
            "probability": round(best_acc * 100, 2),
            "rule_used": rule_str
        }
        
    # Calculate aggregate metrics
    probs = [info["probability"] for info in predictions.values()]
    mean_accuracy = round(np.mean(probs), 2) if probs else 0.0
    median_accuracy = round(np.median(probs), 2) if probs else 0.0
    
    output = {
        "generated_at": datetime.now().isoformat(),
        "forecast_date": forecast_date,
        "mean_accuracy": mean_accuracy,
        "median_accuracy": median_accuracy,
        "predictions": predictions
    }
    
    with open(PREDICTIONS_FILE, "w") as f:
        json.dump(output, f, indent=4)
        
    print(f"Generated predictions for {forecast_date}. Saved to {PREDICTIONS_FILE}")
    print(f"Overall Test Accuracy: {mean_accuracy}%")
    return output

if __name__ == "__main__":
    generate_predictions()