import streamlit as st
import pandas as pd
import numpy as np
import yfinance as yf
import pickle
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import os
from datetime import datetime, timedelta
import warnings
warnings.filterwarnings('ignore')

# Import custom modules
from data_processor import DataProcessor
from technical_indicators import TechnicalIndicators
from model_trainer import ModelTrainer
from database import DatabaseManager

# Configure page
st.set_page_config(
    page_title="AI Stock Trading Signals",
    page_icon="📈",
    layout="wide"
)

# Initialize session state
if 'model_trained' not in st.session_state:
    st.session_state.model_trained = False
if 'model' not in st.session_state:
    st.session_state.model = None
if 'feature_columns' not in st.session_state:
    st.session_state.feature_columns = None
if 'db_manager' not in st.session_state:
    st.session_state.db_manager = DatabaseManager()

def load_or_train_model(stock_symbol):
    """Load existing model or train new one"""
    model_file = f"model_{stock_symbol.replace('.', '_')}.pkl"
    
    if os.path.exists(model_file) and st.session_state.model is None:
        try:
            with open(model_file, 'rb') as f:
                model_data = pickle.load(f)
                st.session_state.model = model_data['model']
                st.session_state.feature_columns = model_data['feature_columns']
                st.session_state.model_trained = True
                return True
        except Exception as e:
            st.error(f"Error loading model: {str(e)}")
            return False
    
    return st.session_state.model_trained

def train_new_model(stock_symbol):
    """Train a new model for the selected stock"""
    with st.spinner("Training AI model... This may take a few minutes."):
        try:
            # Initialize components
            data_processor = DataProcessor()
            tech_indicators = TechnicalIndicators()
            model_trainer = ModelTrainer()
            
            # Download and process data
            data = data_processor.download_stock_data(stock_symbol, period="1y")
            if data is None or len(data) < 100:
                st.error("Insufficient data for training. Please try another stock.")
                return False
            
            # Save stock data to database
            st.session_state.db_manager.save_stock_data(stock_symbol, data)
            
            # Calculate technical indicators
            data_with_indicators = tech_indicators.calculate_all_indicators(data)
            
            # Save technical indicators to database
            st.session_state.db_manager.save_technical_indicators(stock_symbol, data_with_indicators)
            
            # Prepare features and labels
            features, labels = data_processor.prepare_features_and_labels(data_with_indicators)
            
            if len(features) == 0:
                st.error("No valid features could be generated for training.")
                return False
            
            # Train model
            model, feature_columns, metrics = model_trainer.train_model(features, labels)
            
            # Save model performance to database
            st.session_state.db_manager.save_model_metrics(
                stock_symbol, 
                "RandomForest", 
                metrics['accuracy'], 
                metrics['precision'], 
                metrics['recall'],
                len(feature_columns),
                len(features)
            )
            
            # Save model
            model_file = f"model_{stock_symbol.replace('.', '_')}.pkl"
            with open(model_file, 'wb') as f:
                pickle.dump({
                    'model': model,
                    'feature_columns': feature_columns,
                    'metrics': metrics
                }, f)
            
            # Update session state
            st.session_state.model = model
            st.session_state.feature_columns = feature_columns
            st.session_state.model_trained = True
            
            # Display training metrics
            st.success("Model trained successfully!")
            col1, col2, col3 = st.columns(3)
            with col1:
                st.metric("Accuracy", f"{metrics['accuracy']:.3f}")
            with col2:
                st.metric("Precision", f"{metrics['precision']:.3f}")
            with col3:
                st.metric("Recall", f"{metrics['recall']:.3f}")
            
            return True
            
        except Exception as e:
            st.error(f"Error training model: {str(e)}")
            return False

def get_live_price(stock_symbol):
    """Get live stock price"""
    try:
        # Try to get live price using nsepython
        try:
            from nsepython import nsefetch
            # Extract stock name without .NS suffix
            stock_name = stock_symbol.replace('.NS', '')
            url = f"https://www.nseindia.com/api/quote-equity?symbol={stock_name}"
            live_data = nsefetch(url)
            if live_data and 'priceInfo' in live_data and 'lastPrice' in live_data['priceInfo']:
                return float(live_data['priceInfo']['lastPrice'])
        except:
            pass
        
        # Fallback to yfinance
        ticker = yf.Ticker(stock_symbol)
        data = ticker.history(period="1d", interval="1m")
        if not data.empty:
            return float(data['Close'].iloc[-1])
        
        return None
    except Exception as e:
        st.warning(f"Could not fetch live price: {str(e)}")
        return None

def predict_signal(stock_symbol):
    """Predict trading signal for the stock"""
    try:
        # Initialize components
        data_processor = DataProcessor()
        tech_indicators = TechnicalIndicators()
        
        # Get recent data for prediction
        data = data_processor.download_stock_data(stock_symbol, period="3mo")
        if data is None or len(data) < 50:
            return None, "Insufficient data for prediction"
        
        # Calculate technical indicators
        data_with_indicators = tech_indicators.calculate_all_indicators(data)
        
        # Get latest features
        latest_features = data_processor.get_latest_features(data_with_indicators, st.session_state.feature_columns)
        
        if latest_features is None:
            return None, "Could not extract features from latest data"
        
        # Make prediction
        prediction = st.session_state.model.predict([latest_features])[0]
        prediction_proba = st.session_state.model.predict_proba([latest_features])[0]
        
        # Convert prediction to signal
        signal_map = {-1: "SELL", 0: "HOLD", 1: "BUY"}
        signal = signal_map.get(prediction, "HOLD")
        confidence = max(prediction_proba)
        
        # Get current live price for saving to database
        live_price = get_live_price(stock_symbol)
        
        # Save trading signal to database
        if live_price and hasattr(st.session_state, 'db_manager') and st.session_state.db_manager:
            try:
                st.session_state.db_manager.save_trading_signal(
                    stock_symbol, signal, confidence, live_price, model_version="v1.0"
                )
            except Exception as e:
                print(f"Error saving trading signal: {str(e)}")
        
        return {
            'signal': signal,
            'confidence': confidence,
            'prediction': prediction,
            'probabilities': prediction_proba
        }, None
        
    except Exception as e:
        return None, f"Error making prediction: {str(e)}"

def create_price_chart(stock_symbol):
    """Create interactive price chart with technical indicators"""
    try:
        # Initialize components
        data_processor = DataProcessor()
        tech_indicators = TechnicalIndicators()
        
        # Get data for last 30 days
        data = data_processor.download_stock_data(stock_symbol, period="2mo")
        if data is None:
            return None
        
        # Calculate technical indicators
        data_with_indicators = tech_indicators.calculate_all_indicators(data)
        
        # Get last 30 days
        data_chart = data_with_indicators.tail(30).copy()
        
        # Create subplots
        fig = make_subplots(
            rows=4, cols=1,
           # shared_xaxis=True,
            shared_xaxes=True,
            vertical_spacing=0.03,
            subplot_titles=('Price & Bollinger Bands', 'RSI', 'MACD', 'Volume'),
            row_heights=[0.5, 0.2, 0.2, 0.1]
        )
        
        # Price and Bollinger Bands
        fig.add_trace(
            go.Candlestick(
                x=data_chart.index,
                open=data_chart['Open'],
                high=data_chart['High'],
                low=data_chart['Low'],
                close=data_chart['Close'],
                name='Price'
            ), row=1, col=1
        )
        
        # Bollinger Bands
        if 'BB_upper' in data_chart.columns:
            fig.add_trace(
                go.Scatter(x=data_chart.index, y=data_chart['BB_upper'], 
                          name='BB Upper', line=dict(color='red', width=1)), row=1, col=1
            )
            fig.add_trace(
                go.Scatter(x=data_chart.index, y=data_chart['BB_middle'], 
                          name='BB Middle', line=dict(color='blue', width=1)), row=1, col=1
            )
            fig.add_trace(
                go.Scatter(x=data_chart.index, y=data_chart['BB_lower'], 
                          name='BB Lower', line=dict(color='red', width=1)), row=1, col=1
            )
        
        # EMA
        if 'EMA_20' in data_chart.columns:
            fig.add_trace(
                go.Scatter(x=data_chart.index, y=data_chart['EMA_20'], 
                          name='EMA 20', line=dict(color='orange', width=2)), row=1, col=1
            )
        
        # RSI
        if 'RSI' in data_chart.columns:
            fig.add_trace(
                go.Scatter(x=data_chart.index, y=data_chart['RSI'], 
                          name='RSI', line=dict(color='purple')), row=2, col=1
            )
            # RSI overbought/oversold lines
            fig.add_hline(y=70, line_dash="dash", line_color="red", row=2, col=1)
            fig.add_hline(y=30, line_dash="dash", line_color="green", row=2, col=1)
        
        # MACD
        if 'MACD' in data_chart.columns:
            fig.add_trace(
                go.Scatter(x=data_chart.index, y=data_chart['MACD'], 
                          name='MACD', line=dict(color='blue')), row=3, col=1
            )
            if 'MACD_signal' in data_chart.columns:
                fig.add_trace(
                    go.Scatter(x=data_chart.index, y=data_chart['MACD_signal'], 
                              name='MACD Signal', line=dict(color='red')), row=3, col=1
                )
        
        # Volume
        fig.add_trace(
            go.Bar(x=data_chart.index, y=data_chart['Volume'], 
                   name='Volume', marker_color='lightblue'), row=4, col=1
        )
        
        # Update layout
        fig.update_layout(
            title=f"{stock_symbol} - Price Chart with Technical Indicators",
            xaxis_rangeslider_visible=False,
            height=800,
            showlegend=True
        )
        
        # Update y-axis labels
        fig.update_yaxes(title_text="Price (₹)", row=1, col=1)
        fig.update_yaxes(title_text="RSI", row=2, col=1)
        fig.update_yaxes(title_text="MACD", row=3, col=1)
        fig.update_yaxes(title_text="Volume", row=4, col=1)
        
        return fig
        
    except Exception as e:
        st.error(f"Error creating chart: {str(e)}")
        return None

def main():
    st.title("📈 AI-Powered Stock Trading Signals")
    st.markdown("### NSE Stock Analysis with Machine Learning")
    
    # Initialize database manager early
    if 'db_manager' not in st.session_state:
        try:
            from database import DatabaseManager
            st.session_state.db_manager = DatabaseManager()
            st.success("Database initialized successfully")
        except Exception as e:
            st.error(f"Database initialization failed: {str(e)}")
            st.session_state.db_manager = None
    
    # Sidebar for stock selection
    st.sidebar.header("Stock Selection")
    
    # Expanded NSE stock options
    stock_options = {
        "Reliance Industries": "RELIANCE.NS",
        "Tata Consultancy Services": "TCS.NS", 
        "Infosys": "INFY.NS",
        "HDFC Bank": "HDFCBANK.NS",
        "ICICI Bank": "ICICIBANK.NS",
        "State Bank of India": "SBIN.NS",
        "ITC Limited": "ITC.NS",
        "Hindustan Unilever": "HINDUNILVR.NS",
        "Larsen & Toubro": "LT.NS",
        "Axis Bank": "AXISBANK.NS",
        "Kotak Mahindra Bank": "KOTAKBANK.NS",
        "Wipro": "WIPRO.NS",
        "Tech Mahindra": "TECHM.NS",
        "Bharti Airtel": "BHARTIARTL.NS",
        "Asian Paints": "ASIANPAINT.NS",
        "Maruti Suzuki": "MARUTI.NS",
        "Bajaj Finance": "BAJFINANCE.NS",
        "NTPC": "NTPC.NS",
        "Power Grid Corp": "POWERGRID.NS",
        "Tata Steel": "TATASTEEL.NS"
    }
    
    # Option to enter custom stock symbol
    st.sidebar.subheader("Or Enter Custom Stock")
    custom_stock = st.sidebar.text_input(
        "Enter NSE stock symbol (e.g., ADANIPORTS.NS):",
        placeholder="SYMBOL.NS"
    )
    
    # Determine selected stock
    if custom_stock and custom_stock.strip():
        if not custom_stock.endswith('.NS'):
            custom_stock += '.NS'
        selected_stock = custom_stock.upper()
        selected_stock_name = f"Custom: {selected_stock}"
    else:
        selected_stock_name = st.sidebar.selectbox(
            "Choose a stock:",
            list(stock_options.keys())
        )
        selected_stock = stock_options[selected_stock_name]
    
    # Model training section
    st.sidebar.header("Model Management")
    
    # Check if model exists
    model_exists = load_or_train_model(selected_stock)
    
    if not model_exists:
        st.sidebar.warning("No trained model found for this stock.")
        if st.sidebar.button("Train New Model"):
            train_new_model(selected_stock)
    else:
        st.sidebar.success("Model loaded successfully!")
        if st.sidebar.button("Retrain Model"):
            st.session_state.model = None
            st.session_state.model_trained = False
            train_new_model(selected_stock)
    
    # Main content
    if st.session_state.model_trained and st.session_state.model is not None:
        
        # Live price section
        st.header("📊 Live Market Data")
        col1, col2 = st.columns(2)
        
        with col1:
            if st.button("🔄 Get Live Price", type="primary"):
                live_price = get_live_price(selected_stock)
                if live_price:
                    st.metric("Live Price (LTP)", f"₹{live_price:.2f}")
                else:
                    st.warning("Could not fetch live price")
        
        # Prediction section
        st.header("🤖 AI Trading Signal")
        
        if st.button("🎯 Get Trading Signal", type="primary"):
            prediction_result, error = predict_signal(selected_stock)
            
            if prediction_result:
                signal = prediction_result['signal']
                confidence = prediction_result['confidence']
                
                # Display signal with color coding
                if signal == "BUY":
                    st.success(f"🟢 **{signal}** Signal")
                elif signal == "SELL":
                    st.error(f"🔴 **{signal}** Signal")
                else:
                    st.info(f"🟡 **{signal}** Signal")
                
                # Display confidence and probabilities
                col1, col2 = st.columns(2)
                with col1:
                    st.metric("Confidence", f"{confidence:.1%}")
                
                with col2:
                    prob_df = pd.DataFrame({
                        'Signal': ['SELL', 'HOLD', 'BUY'],
                        'Probability': prediction_result['probabilities']
                    })
                    st.dataframe(prob_df, hide_index=True)
                
            else:
                st.error(f"Prediction failed: {error}")
        
        # Chart section
        st.header("📈 Technical Analysis Chart")
        
        if st.button("📊 Generate Chart"):
            chart = create_price_chart(selected_stock)
            if chart:
                st.plotly_chart(chart, use_container_width=True)
        
        # Additional information
        st.header("ℹ️ Model Information")
        
        col1, col2 = st.columns(2)
        with col1:
            st.info("**Features Used:**\n- RSI, EMA, MACD\n- Bollinger Bands\n- SMA, Stochastic\n- Williams %R, ATR")
        
        with col2:
            st.info("**Signal Criteria:**\n- BUY: Expected 2%+ gain\n- SELL: Expected 2%+ loss\n- HOLD: Neutral expectation")
    
    else:
        st.info("👆 Please train a model first using the sidebar to get started with AI-powered trading signals.")
        
        # Show sample data while waiting
        st.header("📊 Sample Stock Data")
        try:
            data_processor = DataProcessor()
            #sample_data = data_processor.download_stock_data(selected_stock, period="1mo")
            sample_data = data_processor.download_stock_data(selected_stock, period="1y")  # 1 year
            # or
            # sample_data = data_processor.download_stock_data(selected_stock, period="2y")  # 2 years
            # or
            # sample_data = data_processor.download_stock_data(selected_stock, period="max")  # full available history
            if sample_data is not None:
                st.dataframe(sample_data.tail(10))
        except:
            st.warning("Could not load sample data")
    
    # Database Information Section
    st.sidebar.header("📊 Database Status")
    
    # Database Console Link
    st.sidebar.markdown("---")
    st.sidebar.markdown("### 🗄️ Database Console")
    st.sidebar.markdown("Access the SQLite database console:")
    st.sidebar.markdown("**[Open Database Console](http://localhost:5001)**")
    st.sidebar.markdown("*View tables, run SQL queries, analyze trading data*")
    st.sidebar.markdown("*Database file: `stock_trading.db`*")
    st.sidebar.markdown("---")
    
    try:
        if hasattr(st.session_state, 'db_manager') and st.session_state.db_manager:
            db_stats = st.session_state.db_manager.get_database_stats()
            if db_stats:
                st.sidebar.metric("Stock Records", db_stats.get('total_stock_records', 0))
                st.sidebar.metric("Trading Signals", db_stats.get('total_signals', 0))
                st.sidebar.metric("Active Models", db_stats.get('active_models', 0))
                st.sidebar.metric("Symbols Tracked", db_stats.get('symbols_tracked', 0))
            else:
                st.sidebar.info("Database ready - no data yet")
        else:
            st.sidebar.info("Database initializing...")
    except Exception as e:
        st.sidebar.warning("Database tables will be created when first used")
    
    # Trading History Section
    if st.session_state.model_trained and hasattr(st.session_state, 'db_manager') and st.session_state.db_manager:
        st.header("📈 Trading History")
        
        col1, col2 = st.columns(2)
        
        with col1:
            if st.button("View Signal History"):
                try:
                    signal_history = st.session_state.db_manager.get_trading_signals_history(selected_stock)
                    if not signal_history.empty:
                        st.dataframe(signal_history.head(10))
                    else:
                        st.info("No trading signals found in database")
                except Exception as e:
                    st.error(f"Error loading signal history: {str(e)}")
        
        with col2:
            if st.button("Model Performance"):
                try:
                    model_perf = st.session_state.db_manager.get_model_performance(selected_stock)
                    if model_perf:
                        st.json(model_perf)
                    else:
                        st.info("No model performance data found")
                except Exception as e:
                    st.error(f"Error loading model performance: {str(e)}")
                    
        # Database Console Access
        st.markdown("---")
        st.info("💡 **Tip**: Use the Database Console in the sidebar to view all stored data, run SQL queries, and analyze trading patterns across all stocks.")

if __name__ == "__main__":
    main()