import yfinance as yf
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from sklearn.preprocessing import MinMaxScaler
from datetime import datetime, timedelta

def get_idx_stocks():
    """Return dictionary of major IDX stocks"""
    return {
        "BBCA.JK": "Bank Central Asia",
        "BMRI.JK": "Bank Mandiri",
        "BBNI.JK": "Bank BNI",
        "BBRI.JK": "Bank BRI",
        "TLKM.JK": "Telkom Indonesia",
        "UNVR.JK": "Unilever Indonesia",
        "INDF.JK": "Indofood Sukses Makmur",
        "KLBF.JK": "Kalbe Farma",
        "ASII.JK": "Astra International",
        "ICBP.JK": "Indofood CBP Sukses Makmur",
        "SMGR.JK": "Semen Indonesia",
        "INTP.JK": "Indocement Tunggal Prakasa",
        "ANTM.JK": "Aneka Tambang",
        "TINS.JK": "Timah",
        "PTBA.JK": "Bukit Asam",
        "PGAS.JK": "Perusahaan Gas Negara",
        "EXCL.JK": "XL Axiata",
        "ISAT.JK": "Indosat Ooredoo",
        "FREN.JK": "Smartfren Telecom",
        "BKLA.JK": "Bank Bukopin",
        "BANK.JK": "Bank Danamon",
        "BSIM.JK": "Bank Syariah Indonesia",
        "MDKA.JK": "Merdeka Copper Gold",
        "SMCB.JK": "Semen Baturaja",
        "WIKA.JK": "Wijaya Karya",
        "ADHI.JK": "Adhi Karya",
        "PTPP.JK": "PP (Persero)",
        "JSMR.JK": "Jasa Marga",
        "TPIA.JK": "Chandra Asri Petrochemical",
        "SRIL.JK": "Sri Rejeki Isman",
        "GGRM.JK": "Gudang Garam",
        "HMSP.JK": "HM Sampoerna",
        "TCID.JK": "MNC Investama",
        "MNCN.JK": "Media Nusantara Citra",
        "BHITT.JK": "Bumi Resources Minerals",
        "DOID.JK": "Delta Dunia Makmur",
        "MEDC.JK": "Medco Energi International",
        "PGAS.JK": "Perusahaan Gas Negara"
    }

def fetch_stock_data(symbol, period_days):
    """Fetch stock data from Yahoo Finance"""
    try:
        # Calculate end date (today) and start date
        end_date = datetime.now()
        start_date = end_date - timedelta(days=period_days + 30)  # Extra days to ensure we get enough trading days
        
        # Download data
        stock = yf.Ticker(symbol)
        data = stock.history(start=start_date, end=end_date)
        
        if data.empty:
            return None
        
        # Take only the requested period (accounting for weekends)
        trading_days = 0
        filtered_data = []
        
        for i in range(len(data)-1, -1, -1):
            filtered_data.append(data.iloc[i])
            trading_days += 1
            if trading_days >= period_days:
                break
        
        filtered_data.reverse()
        result = pd.DataFrame(filtered_data)
        result.index = pd.to_datetime(result.index)
        
        return result
        
    except Exception as e:
        print(f"Error fetching data for {symbol}: {e}")
        return None

def prepare_timesfm_data(stock_data, use_volume=False):
    """Prepare stock data for TimesFM model"""
    # Select features
    if use_volume:
        features = ['Close', 'Volume']
    else:
        features = ['Close']
    
    data = stock_data[features].copy()
    
    # Handle missing values
    data = data.fillna(method='ffill').fillna(method='bfill')
    
    # Scale the data
    scaler = MinMaxScaler()
    scaled_data = scaler.fit_transform(data)
    
    # For TimesFM, we'll use the closing prices (and optionally volume)
    if use_volume:
        # Combine close and volume into a single series
        # Here we'll use closing prices as primary and volume as secondary
        timesfm_data = scaled_data[:, 0]  # Use closing prices
    else:
        timesfm_data = scaled_data.flatten()
    
    return timesfm_data, scaler

def calculate_metrics(actual, predicted):
    """Calculate evaluation metrics"""
    mae = np.mean(np.abs(actual - predicted))
    mse = np.mean((actual - predicted) ** 2)
    rmse = np.sqrt(mse)
    mape = np.mean(np.abs((actual - predicted) / actual)) * 100
    
    return {
        'MAE': mae,
        'MSE': mse,
        'RMSE': rmse,
        'MAPE': mape
    }

def create_forecast_plot(historical_data, forecast_dates, forecast_prices, symbol):
    """Create an interactive plot with historical data and forecast"""
    fig = make_subplots(
        rows=2, cols=1,
        shared_xaxes=True,
        vertical_spacing=0.05,
        subplot_titles=(f'{symbol} Stock Price Forecast', 'Volume'),
        row_width=[0.2, 0.7]
    )
    
    # Historical closing prices
    fig.add_trace(
        go.Scatter(
            x=historical_data.index,
            y=historical_data['Close'],
            mode='lines',
            name='Historical Price',
            line=dict(color='blue', width=2)
        ),
        row=1, col=1
    )
    
    # Forecast prices
    fig.add_trace(
        go.Scatter(
            x=forecast_dates,
            y=forecast_prices,
            mode='lines',
            name='Forecast',
            line=dict(color='red', width=2, dash='dash')
        ),
        row=1, col=1
    )
    
    # Confidence interval (simple approach)
    std_dev = np.std(historical_data['Close'].values) * 0.1
    upper_bound = forecast_prices + std_dev
    lower_bound = forecast_prices - std_dev
    
    fig.add_trace(
        go.Scatter(
            x=forecast_dates,
            y=upper_bound,
            mode='lines',
            line=dict(width=0),
            showlegend=False,
            hoverinfo='skip'
        ),
        row=1, col=1
    )
    
    fig.add_trace(
        go.Scatter(
            x=forecast_dates,
            y=lower_bound,
            mode='lines',
            line=dict(width=0),
            fill='tonexty',
            fillcolor='rgba(255,0,0,0.2)',
            name='Confidence Interval',
            hoverinfo='skip'
        ),
        row=1, col=1
    )
    
    # Volume
    fig.add_trace(
        go.Bar(
            x=historical_data.index,
            y=historical_data['Volume'],
            name='Volume',
            marker_color='lightblue',
            yaxis='y2'
        ),
        row=2, col=1
    )
    
    # Update layout
    fig.update_layout(
        title=f'{symbol} Stock Price Prediction',
        xaxis_title='Date',
        yaxis_title='Price (IDR)',
        yaxis2_title='Volume',
        hovermode='x unified',
        showlegend=True,
        height=600
    )
    
    fig.update_yaxes(title_text="Price (IDR)", row=1, col=1)
    fig.update_yaxes(title_text="Volume", row=2, col=1)
    fig.update_xaxes(title_text="Date", row=2, col=1)
    
    return fig

def get_stock_info(symbol):
    """Get additional stock information"""
    try:
        stock = yf.Ticker(symbol)
        info = stock.info
        
        # Format market cap
        market_cap = info.get('marketCap', 0)
        if market_cap:
            if market_cap >= 1e12:
                market_cap_str = f"Rp {market_cap/1e12:.2f}T"
            elif market_cap >= 1e9:
                market_cap_str = f"Rp {market_cap/1e9:.2f}B"
            elif market_cap >= 1e6:
                market_cap_str = f"Rp {market_cap/1e6:.2f}M"
            else:
                market_cap_str = f"Rp {market_cap:,.0f}"
            info['marketCap'] = market_cap_str
        
        return info
    except Exception as e:
        print(f"Error getting stock info: {e}")
        return {}