import gradio as gr
import os
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from datetime import datetime, timedelta
import io
import base64
from dotenv import load_dotenv
import json
import subprocess
import sys
from typing import Optional, Dict, Any

# Load environment variables from .env file
load_dotenv()

# Import financial and analysis libraries
try:
    import yfinance as yf
    YFINANCE_AVAILABLE = True
except ImportError:
    YFINANCE_AVAILABLE = False

try:
    from scipy import stats
    import scipy.optimize as optimize
    SCIPY_AVAILABLE = True
except ImportError:
    SCIPY_AVAILABLE = False

try:
    from sklearn.linear_model import LinearRegression
    from sklearn.preprocessing import StandardScaler
    SKLEARN_AVAILABLE = True
except ImportError:
    SKLEARN_AVAILABLE = False

try:
    import seaborn as sns
    import matplotlib.pyplot as plt
    VISUALIZATION_AVAILABLE = True
except ImportError:
    VISUALIZATION_AVAILABLE = False
    plt = None

# Try to import OpenAI
try:
    from openai import OpenAI
    openai_client = None
    api_key = os.getenv("OPENAI_API_KEY")
    if api_key:
        openai_client = OpenAI(api_key=api_key)
        print("OpenAI client initialized successfully")
    else:
        print("No OpenAI API key found")
except ImportError:
    openai_client = None
    print("OpenAI library not available")

# Financial calculation functions
def calculate_dcf(free_cash_flows, terminal_growth_rate=0.02, wacc=0.10):
    """Calculate DCF valuation."""
    try:
        present_values = []
        for i, fcf in enumerate(free_cash_flows):
            pv = fcf / ((1 + wacc) ** (i + 1))
            present_values.append(pv)

        # Terminal value
        terminal_fcf = free_cash_flows[-1] * (1 + terminal_growth_rate)
        terminal_value = terminal_fcf / (wacc - terminal_growth_rate)
        terminal_pv = terminal_value / ((1 + wacc) ** len(free_cash_flows))

        total_dcf = sum(present_values) + terminal_pv

        return {
            "dcf_value": total_dcf,
            "present_values": present_values,
            "terminal_value": terminal_value,
            "terminal_pv": terminal_pv
        }
    except Exception as e:
        return {"error": str(e)}

def calculate_sharpe_ratio(returns, risk_free_rate=0.02):
    """Calculate Sharpe ratio."""
    try:
        excess_returns = np.mean(returns) - risk_free_rate/252  # Daily risk-free rate
        volatility = np.std(returns)
        sharpe = excess_returns / volatility * np.sqrt(252)  # Annualized
        return sharpe
    except Exception as e:
        return f"Error: {str(e)}"

def get_stock_data(symbol, period="1y"):
    """Fetch stock data using yfinance."""
    if not YFINANCE_AVAILABLE:
        return None

    try:
        stock = yf.Ticker(symbol)
        data = stock.history(period=period)
        info = stock.info
        return {"data": data, "info": info}
    except Exception as e:
        return {"error": str(e)}

def create_stock_chart(symbol, data):
    """Create interactive stock chart with enhanced styling."""
    try:
        fig = go.Figure()

        # Add candlestick chart
        fig.add_trace(go.Candlestick(
            x=data.index,
            open=data['Open'],
            high=data['High'],
            low=data['Low'],
            close=data['Close'],
            name=symbol,
            increasing_line_color='#00C851',
            decreasing_line_color='#FF4444'
        ))

        # Add moving averages
        ma20 = data['Close'].rolling(window=20).mean()
        ma50 = data['Close'].rolling(window=50).mean()

        fig.add_trace(go.Scatter(
            x=data.index, y=ma20, name='MA20',
            line=dict(color='#007bff', width=1.5)
        ))

        fig.add_trace(go.Scatter(
            x=data.index, y=ma50, name='MA50',
            line=dict(color='#6f42c1', width=1.5)
        ))

        # Enhanced layout with dark theme
        fig.update_layout(
            title={
                'text': f"📈 {symbol} Stock Analysis",
                'x': 0.5,
                'font': {'size': 24, 'color': '#2c3e50'}
            },
            yaxis_title="Price ($)",
            xaxis_title="Date",
            template="plotly_white",
            height=600,
            showlegend=True,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            ),
            margin=dict(l=50, r=50, t=80, b=50)
        )

        return fig
    except Exception as e:
        return f"Chart error: {str(e)}"

def perform_financial_analysis(symbol):
    """Comprehensive financial analysis."""
    if not YFINANCE_AVAILABLE:
        return "Install yfinance for real-time analysis: pip install yfinance"

    try:
        stock_data = get_stock_data(symbol)
        if "error" in stock_data:
            return f"Error fetching data: {stock_data['error']}"

        data = stock_data["data"]
        info = stock_data["info"]

        # Calculate returns
        returns = data['Close'].pct_change().dropna()

        # Calculate metrics
        sharpe = calculate_sharpe_ratio(returns)
        volatility = returns.std() * np.sqrt(252)  # Annualized

        # Beta calculation (vs SPY)
        try:
            spy_data = get_stock_data("SPY", period="1y")["data"]
            spy_returns = spy_data['Close'].pct_change().dropna()

            # Align dates
            aligned_data = pd.concat([returns, spy_returns], axis=1, keys=[symbol, 'SPY']).dropna()

            if SKLEARN_AVAILABLE:
                model = LinearRegression()
                X = aligned_data['SPY'].values.reshape(-1, 1)
                y = aligned_data[symbol].values
                model.fit(X, y)
                beta = model.coef_[0]
            else:
                beta = aligned_data.cov().iloc[0, 1] / aligned_data['SPY'].var()
        except:
            beta = "N/A"

        # Current metrics
        current_price = data['Close'][-1]
        pe_ratio = info.get('trailingPE', 'N/A')
        market_cap = info.get('marketCap', 'N/A')

        analysis = f"""
## {symbol} Financial Analysis

### **Current Metrics**
- **Price**: ${current_price:.2f}
- **Market Cap**: {f"${market_cap:,.0f}" if market_cap != 'N/A' else 'N/A'}
- **P/E Ratio**: {pe_ratio}

### **Risk Metrics**
- **Beta**: {f"{beta:.2f}" if beta != "N/A" else "N/A"}
- **Sharpe Ratio**: {sharpe:.3f}
- **Volatility**: {volatility:.1%}

### **Performance**
- **1-Month Return**: {returns.tail(21).sum():.1%}
- **3-Month Return**: {returns.tail(63).sum():.1%}
- **YTD Return**: {returns.sum():.1%}

*Analysis powered by yfinance, numpy, and pandas*
        """

        return analysis.strip()

    except Exception as e:
        return f"Analysis error: {str(e)}"

def execute_python_code(code: str) -> str:
    """Execute Python code safely and return results."""
    try:
        # Create a safe environment with limited imports
        safe_globals = {
            '__builtins__': {
                'print': print, 'len': len, 'range': range, 'enumerate': enumerate,
                'sum': sum, 'min': min, 'max': max, 'abs': abs, 'round': round,
                'str': str, 'int': int, 'float': float, 'bool': bool, 'list': list,
                'dict': dict, 'tuple': tuple, 'set': set
            },
            'np': np, 'pd': pd, 'plt': plt if VISUALIZATION_AVAILABLE else None,
            'yf': yf if YFINANCE_AVAILABLE else None
        }

        # Capture output
        old_stdout = sys.stdout
        sys.stdout = captured_output = io.StringIO()

        try:
            exec(code, safe_globals)
            result = captured_output.getvalue()
            return result if result else "Code executed successfully (no output)"
        finally:
            sys.stdout = old_stdout

    except Exception as e:
        return f"Error: {str(e)}"

def generate_financial_visualization(prompt: str) -> Optional[str]:
    """Generate financial visualization using DALL-E."""
    if not openai_client:
        print("No OpenAI client available")
        return None

    try:
        print(f"Generating image with prompt: {prompt}")

        # Create a more specific financial prompt
        enhanced_prompt = f"Professional financial data visualization: {prompt}. Clean, modern design with charts, graphs, and financial elements. Corporate style with blue and white colors."

        response = openai_client.images.generate(
            model="dall-e-3",
            prompt=enhanced_prompt,
            size="1024x1024",
            quality="standard",
            n=1
        )

        image_url = response.data[0].url
        print(f"Successfully generated image: {image_url}")
        return image_url

    except Exception as e:
        print(f"DALL-E error details: {type(e).__name__}: {str(e)}")
        # Try to get more specific error information
        if hasattr(e, 'response'):
            print(f"Response status: {e.response.status_code if hasattr(e.response, 'status_code') else 'unknown'}")
        return None

def get_openai_code_response(message: str) -> Optional[str]:
    """Get code-focused response from OpenAI."""
    if not openai_client:
        return None

    try:
        system_prompt = """You are an expert Python developer specializing in financial analysis and data science.

Provide working Python code examples using these libraries:
- yfinance for stock data
- pandas/numpy for data analysis
- matplotlib/plotly for visualizations
- scipy for statistical analysis
- scikit-learn for machine learning

Always include:
1. Working code examples
2. Clear comments explaining each step
3. Error handling where appropriate
4. Practical financial applications

Format responses with proper markdown code blocks."""

        response = openai_client.chat.completions.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": message}
            ],
            max_tokens=1500,
            temperature=0.1
        )
        return response.choices[0].message.content
    except Exception as e:
        print(f"OpenAI code error: {e}")
        return None

def get_openai_response(message):
    """Get enhanced response from OpenAI with financial context."""
    if not openai_client:
        return None

    try:
        # Enhanced system prompt with library context
        system_prompt = f"""You are CFAx Agent, an expert financial AI assistant with access to powerful Python libraries for financial analysis.

Available tools and libraries:
- yfinance: {'✅' if YFINANCE_AVAILABLE else '❌'} (real-time stock data)
- scipy: {'✅' if SCIPY_AVAILABLE else '❌'} (statistical analysis, optimization)
- scikit-learn: {'✅' if SKLEARN_AVAILABLE else '❌'} (machine learning, regression analysis)
- plotly/matplotlib: {'✅' if VISUALIZATION_AVAILABLE else '❌'} (advanced visualizations)
- pandas/numpy: ✅ (data analysis and calculations)

Provide professional financial analysis with:
1. Clear explanations suitable for CFA-level analysis
2. Specific calculations and formulas when relevant
3. Risk assessment and portfolio theory insights
4. Market context and economic considerations
5. Recommendations based on quantitative analysis

Use markdown formatting for better readability. When discussing stocks or calculations, mention which tools would be used for analysis."""

        response = openai_client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": message}
            ],
            max_tokens=1200,
            temperature=0.3
        )
        return response.choices[0].message.content
    except Exception as e:
        print(f"OpenAI error: {e}")
        return None

def financial_chat(message, history):
    """Enhanced chat function with comprehensive financial analysis."""

    # Check for image generation requests
    if any(word in message.lower() for word in ["image", "visualization", "chart", "generate", "create", "visual", "picture", "diagram"]):
        if openai_client:
            print(f"Processing image generation request: {message}")
            image_url = generate_financial_visualization(message)
            if image_url:
                return f"🎨 **Generated Financial Visualization**\n\n![Generated Image]({image_url})\n\n*AI-generated visualization based on your request: \"{message}\"*"
            else:
                return "❌ **Image Generation Failed** - There was an issue generating the visualization. This could be due to:\n\n• API rate limits\n• Content policy restrictions\n• Network connectivity issues\n\nPlease try:\n• A simpler request like 'financial dashboard'\n• Waiting a moment and trying again\n• Using different keywords"
        else:
            return "🔑 **OpenAI API Key Required** - Add your OpenAI API key to enable AI-generated visualizations"

    # Check for specific analysis requests
    if any(word in message.lower() for word in ["analyze", "analysis"]) and any(word in message.upper() for word in ["AAPL", "TSLA", "MSFT", "GOOGL", "AMZN", "NVDA", "META"]):
        # Extract stock symbol
        words = message.upper().split()
        symbols = ["AAPL", "TSLA", "MSFT", "GOOGL", "AMZN", "NVDA", "META"]
        for word in words:
            if word in symbols:
                return perform_financial_analysis(word)

    # Try OpenAI first if available
    if openai_client:
        openai_response = get_openai_response(message)
        if openai_response:
            return f"🤖 **CFAx Agent (Enhanced with OpenAI)**\n\n{openai_response}"

    # Enhanced fallback knowledge base
    msg_lower = message.lower()

    if "libraries" in msg_lower or "tools" in msg_lower:
        return f"""🔧 **Available Financial Analysis Tools**

### **Core Libraries Status:**
- **yfinance**: {'✅ Available' if YFINANCE_AVAILABLE else '❌ Not installed'} - Real-time stock data
- **scipy**: {'✅ Available' if SCIPY_AVAILABLE else '❌ Not installed'} - Statistical analysis, optimization
- **scikit-learn**: {'✅ Available' if SKLEARN_AVAILABLE else '❌ Not installed'} - Machine learning, regression
- **plotly/seaborn**: {'✅ Available' if VISUALIZATION_AVAILABLE else '❌ Not installed'} - Advanced visualizations
- **pandas/numpy**: ✅ Available - Data manipulation and calculations

### **Capabilities:**
- Real-time stock analysis and charting
- DCF valuations with custom inputs
- Risk metrics (Beta, Sharpe ratio, VaR)
- Portfolio optimization
- Statistical modeling and regression analysis
- Interactive financial visualizations

*Install missing libraries: `pip install yfinance scipy scikit-learn matplotlib seaborn`*"""

    elif "apple" in msg_lower or "aapl" in msg_lower:
        if YFINANCE_AVAILABLE:
            return perform_financial_analysis("AAPL")
        else:
            return """📊 **Apple Inc. (AAPL) Analysis**

Apple is a large-cap technology stock with strong fundamentals:
• Market Cap: ~$3 Trillion
• P/E Ratio: ~25-30
• Sector: Technology/Consumer Electronics
• Strong cash flow and dividend paying stock

*Install yfinance for real-time analysis: `pip install yfinance`*"""

    elif any(word in msg_lower for word in ["dcf", "valuation"]):
        return """💰 **DCF Valuation Method**

**Enhanced with Python:**
```python
def calculate_dcf(fcf_list, wacc=0.10, terminal_growth=0.02):
    # Present value of cash flows
    pv_sum = sum(fcf/(1+wacc)**i for i, fcf in enumerate(fcf_list, 1))

    # Terminal value
    terminal_fcf = fcf_list[-1] * (1 + terminal_growth)
    terminal_value = terminal_fcf / (wacc - terminal_growth)
    terminal_pv = terminal_value / (1 + wacc)**len(fcf_list)

    return pv_sum + terminal_pv
```

**Libraries used:** numpy for calculations, pandas for data handling

*Try: "Calculate DCF for [cash flows]" for live calculations*"""

    elif "portfolio" in msg_lower or "optimization" in msg_lower:
        return f"""📈 **Portfolio Optimization**

**Modern Portfolio Theory with Python:**

{'✅ **scipy.optimize**: Efficient frontier calculation' if SCIPY_AVAILABLE else '❌ Install scipy for optimization'}
{'✅ **sklearn**: Risk factor modeling' if SKLEARN_AVAILABLE else '❌ Install scikit-learn for ML models'}
{'✅ **numpy**: Covariance matrices and returns' if True else ''}

**Key Functions:**
- Mean-variance optimization
- Sharpe ratio maximization
- Risk parity portfolios
- Monte Carlo simulations

**Example:**
```python
from scipy.optimize import minimize
# Minimize portfolio risk for target return
weights = minimize(portfolio_risk, initial_weights,
                  constraints=constraints)
```"""

    else:
        api_status = "🟢 Connected" if openai_client else "🔴 Not Connected"
        lib_count = sum([YFINANCE_AVAILABLE, SCIPY_AVAILABLE, SKLEARN_AVAILABLE, VISUALIZATION_AVAILABLE]) + 2  # Always have pandas/numpy

        return f"""**CFAx Agent** - *Your AI Financial Analyst* (OpenAI: {api_status})

**Market-Ready Analytics Platform: {lib_count}/6 Libraries Active**

**What can I help you analyze today?**

📈 **Equity Research**: *"Analyze AAPL's risk-adjusted returns vs benchmark"*
💰 **Valuation Models**: *"Run a DCF analysis on these cash flows: [100, 110, 120]M with 8% WACC"*
🎯 **Portfolio Strategy**: *"Optimize allocation for tech-heavy portfolio with 15% target return"*
📊 **Risk Assessment**: *"Compare volatility profiles between TSLA and traditional auto stocks"*
🔍 **Market Intelligence**: *"What's driving the current sector rotation into value?"*

**Live Market Capabilities:**
Real-time equity data • Monte Carlo simulations • Efficient frontier optimization • Beta calculations • Correlation analysis

{f"🔥 **All systems operational** - Enhanced AI with {lib_count} quantitative libraries at your service!" if openai_client else f"⚡ **Ready for analysis** - Add OpenAI key for enhanced market intelligence!"}

*Ask me anything about markets, valuations, or portfolio strategy - I'm here to deliver institutional-grade financial analysis.*"""

# Enhanced UI Components
def create_stock_analysis_interface():
    """Create the stock analysis tab interface."""
    with gr.Row():
        with gr.Column(scale=1):
            stock_input = gr.Textbox(
                label="Stock Symbol",
                placeholder="Enter symbol (e.g., AAPL, TSLA, MSFT)",
                value="AAPL"
            )
            period_select = gr.Dropdown(
                choices=["1mo", "3mo", "6mo", "1y", "2y", "5y"],
                value="1y",
                label="Time Period"
            )
            analyze_btn = gr.Button(
                "Analyze Stock",
                variant="primary",
                size="lg"
            )

        with gr.Column(scale=2):
            stock_chart = gr.Plot(label="Stock Chart")

    with gr.Row():
        stock_analysis = gr.Markdown(label="Analysis Results")

    def analyze_stock(symbol, period):
        if not symbol:
            return None, "Please enter a stock symbol"

        stock_data = get_stock_data(symbol.upper(), period)
        if not stock_data or "error" in stock_data:
            return None, f"Error fetching data for {symbol}"

        chart = create_stock_chart(symbol.upper(), stock_data["data"])
        analysis = perform_financial_analysis(symbol.upper())

        return chart, analysis

    analyze_btn.click(
        analyze_stock,
        inputs=[stock_input, period_select],
        outputs=[stock_chart, stock_analysis]
    )

    return stock_input, period_select, analyze_btn, stock_chart, stock_analysis

def create_code_executor_interface():
    """Create the code executor tab interface."""
    with gr.Row():
        with gr.Column():
            code_input = gr.Code(
                label="Python Code Editor",
                language="python",
                value="""# Financial Analysis Example
import numpy as np
import pandas as pd

# Calculate simple moving average
prices = [100, 102, 101, 103, 105, 104, 106]
ma_5 = sum(prices[-5:]) / 5
print(f"5-day Moving Average: ${ma_5:.2f}")

# Calculate returns
returns = [(prices[i] - prices[i-1]) / prices[i-1] for i in range(1, len(prices))]
print(f"Daily Returns: {[f'{r:.2%}' for r in returns]}")"""
            )

        with gr.Column():
            code_output = gr.Textbox(
                label="Output",
                lines=10,
                max_lines=20
            )

    with gr.Row():
        execute_btn = gr.Button(
            "Execute Code",
            variant="primary",
            size="lg"
        )
        clear_btn = gr.Button("Clear", variant="secondary")

    with gr.Row():
        code_examples = gr.Dropdown(
            choices=[
                "DCF Calculation",
                "Sharpe Ratio Analysis",
                "Portfolio Optimization",
                "Monte Carlo Simulation",
                "Stock Correlation Analysis"
            ],
            label="Code Examples",
            value=None
        )

    def load_example(example):
        examples = {
            "DCF Calculation": """# DCF Valuation Calculator
free_cash_flows = [100, 110, 121, 133, 146]  # millions
wacc = 0.10  # 10%
terminal_growth = 0.02  # 2%

# Calculate present values
present_values = []
for i, fcf in enumerate(free_cash_flows):
    pv = fcf / ((1 + wacc) ** (i + 1))
    present_values.append(pv)
    print(f"Year {i+1}: FCF ${fcf}M -> PV ${pv:.1f}M")

# Terminal value
terminal_fcf = free_cash_flows[-1] * (1 + terminal_growth)
terminal_value = terminal_fcf / (wacc - terminal_growth)
terminal_pv = terminal_value / ((1 + wacc) ** len(free_cash_flows))

total_dcf = sum(present_values) + terminal_pv
print(f"\nTotal DCF Value: ${total_dcf:.1f}M")
print(f"Terminal Value: ${terminal_pv:.1f}M")""",

            "Sharpe Ratio Analysis": """# Sharpe Ratio Calculation
import numpy as np

# Sample daily returns (as decimals)
returns = np.array([0.01, -0.005, 0.02, 0.015, -0.01, 0.008, 0.012])
risk_free_rate = 0.02 / 252  # 2% annual, daily

# Calculate metrics
mean_return = np.mean(returns)
volatility = np.std(returns)
excess_return = mean_return - risk_free_rate

# Annualized Sharpe ratio
sharpe_ratio = (excess_return / volatility) * np.sqrt(252)

print(f"Average Daily Return: {mean_return:.4f} ({mean_return*252:.2%} annual)")
print(f"Daily Volatility: {volatility:.4f} ({volatility*np.sqrt(252):.2%} annual)")
print(f"Sharpe Ratio: {sharpe_ratio:.3f}")

if sharpe_ratio > 1:
    print("📈 Excellent risk-adjusted performance!")
elif sharpe_ratio > 0.5:
    print("📊 Good risk-adjusted performance")
else:
    print("📉 Poor risk-adjusted performance")""",

            "Portfolio Optimization": """# Simple Portfolio Optimization
import numpy as np

# Asset expected returns and volatilities
assets = ['Stock A', 'Stock B', 'Bond']
expected_returns = np.array([0.12, 0.15, 0.05])  # 12%, 15%, 5%
volatilities = np.array([0.20, 0.25, 0.03])      # 20%, 25%, 3%

# Correlation matrix
correlations = np.array([
    [1.0, 0.3, 0.1],
    [0.3, 1.0, 0.05],
    [0.1, 0.05, 1.0]
])

# Calculate covariance matrix
cov_matrix = np.outer(volatilities, volatilities) * correlations

print("Portfolio Analysis:")
print("==================")
for i, asset in enumerate(assets):
    print(f"{asset}: {expected_returns[i]:.1%} return, {volatilities[i]:.1%} volatility")

# Equal weight portfolio
weights = np.array([1/3, 1/3, 1/3])
portfolio_return = np.dot(weights, expected_returns)
portfolio_risk = np.sqrt(np.dot(weights, np.dot(cov_matrix, weights)))

print(f"\nEqual Weight Portfolio:")
print(f"Expected Return: {portfolio_return:.2%}")
print(f"Portfolio Risk: {portfolio_risk:.2%}")
print(f"Risk-Return Ratio: {portfolio_return/portfolio_risk:.2f}")""",

            "Monte Carlo Simulation": """# Monte Carlo Stock Price Simulation
import numpy as np

# Parameters
initial_price = 100
mu = 0.08  # 8% annual drift
sigma = 0.20  # 20% annual volatility
days = 252  # 1 year
simulations = 1000

# Time step
dt = 1/252

# Generate random price paths
np.random.seed(42)
final_prices = []

for _ in range(simulations):
    price = initial_price
    for day in range(days):
        random_shock = np.random.normal(0, 1)
        price_change = price * (mu * dt + sigma * np.sqrt(dt) * random_shock)
        price += price_change
    final_prices.append(price)

# Analysis
final_prices = np.array(final_prices)
mean_price = np.mean(final_prices)
std_price = np.std(final_prices)

print(f"Monte Carlo Simulation Results ({simulations:,} simulations):")
print(f"Initial Price: ${initial_price:.2f}")
print(f"Mean Final Price: ${mean_price:.2f}")
print(f"Standard Deviation: ${std_price:.2f}")
print(f"Min Price: ${np.min(final_prices):.2f}")
print(f"Max Price: ${np.max(final_prices):.2f}")

# Probability analysis
prob_profit = np.mean(final_prices > initial_price)
prob_loss_10 = np.mean(final_prices < initial_price * 0.9)

print(f"\nProbability of Profit: {prob_profit:.1%}")
print(f"Probability of >10% Loss: {prob_loss_10:.1%}")""",

            "Stock Correlation Analysis": """# Stock Correlation Analysis
import numpy as np

# Sample price data for 3 stocks over 10 days
stock_a = [100, 102, 101, 103, 105, 104, 106, 108, 107, 109]
stock_b = [50, 51, 50.5, 52, 53, 52.5, 54, 55, 54.5, 56]
stock_c = [200, 198, 201, 199, 202, 200, 203, 201, 204, 202]

# Calculate daily returns
def calculate_returns(prices):
    return [(prices[i] - prices[i-1]) / prices[i-1] for i in range(1, len(prices))]

returns_a = calculate_returns(stock_a)
returns_b = calculate_returns(stock_b)
returns_c = calculate_returns(stock_c)

print("Daily Returns Analysis:")
print("=======================")
print(f"Stock A returns: {[f'{r:.2%}' for r in returns_a]}")
print(f"Stock B returns: {[f'{r:.2%}' for r in returns_b]}")
print(f"Stock C returns: {[f'{r:.2%}' for r in returns_c]}")

# Calculate correlation matrix
returns_matrix = np.array([returns_a, returns_b, returns_c])
correlation_matrix = np.corrcoef(returns_matrix)

print("\nCorrelation Matrix:")
print("==================")
stocks = ['Stock A', 'Stock B', 'Stock C']
for i, stock1 in enumerate(stocks):
    for j, stock2 in enumerate(stocks):
        corr = correlation_matrix[i, j]
        print(f"{stock1} vs {stock2}: {corr:.3f}")

# Interpretation
print("\nInterpretation:")
print("===============")
ab_corr = correlation_matrix[0, 1]
if ab_corr > 0.7:
    print(f"Stock A & B are highly correlated ({ab_corr:.3f})")
elif ab_corr > 0.3:
    print(f"Stock A & B are moderately correlated ({ab_corr:.3f})")
else:
    print(f"Stock A & B have low correlation ({ab_corr:.3f})")"""
        }
        return examples.get(example, "")

    def execute_code(code):
        if not code.strip():
            return "Please enter some code to execute."
        return execute_python_code(code)

    def clear_code():
        return "", ""

    execute_btn.click(execute_code, inputs=[code_input], outputs=[code_output])
    clear_btn.click(clear_code, outputs=[code_input, code_output])
    code_examples.change(load_example, inputs=[code_examples], outputs=[code_input])

    return code_input, code_output, execute_btn, clear_btn, code_examples

def create_ai_visualization_interface():
    """Create the AI visualization tab interface."""
    with gr.Row():
        with gr.Column(scale=1):
            viz_prompt = gr.Textbox(
                label="Visualization Prompt",
                placeholder="Describe the financial chart or visualization you want...",
                lines=3
            )
            generate_btn = gr.Button(
                "Generate Visualization",
                variant="primary",
                size="lg"
            )

        with gr.Column(scale=2):
            viz_output = gr.Image(label="Generated Visualization")

    with gr.Row():
        viz_examples = gr.Examples(
            examples=[
                ["Stock market bull vs bear visualization with charts"],
                ["Modern portfolio dashboard with pie charts and graphs"],
                ["Risk-return scatter plot with efficient frontier"],
                ["Financial technology concept with AI and data"],
                ["Cryptocurrency trading interface design"]
            ],
            inputs=[viz_prompt]
        )

    def generate_viz(prompt):
        if not prompt.strip():
            return None

        if not openai_client:
            return None

        image_url = generate_financial_visualization(prompt)
        return image_url

    generate_btn.click(generate_viz, inputs=[viz_prompt], outputs=[viz_output])

    return viz_prompt, generate_btn, viz_output, viz_examples

# Create the main interface with enhanced tabs
with gr.Blocks(
    theme=gr.themes.Soft(),
    title="CFAx Agent - Professional Financial Analysis",
    css="""
    .gradio-container {
        background: linear-gradient(135deg, #6c757d 0%, #495057 100%);
        font-family: 'Segoe UI', system-ui, sans-serif;
        padding-top: 20px !important;
        margin-top: 0 !important;
    }
    .main-container {
        max-width: 1200px;
        margin: 0 auto;
        padding: 10px;
    }
    .header-section {
        margin-top: 0 !important;
        padding-top: 0 !important;
        position: relative;
        z-index: 100;
    }
    .gr-box {
        border-radius: 12px;
        border: 1px solid #e1e5e9;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.07);
    }
    .gr-button {
        border-radius: 8px;
        font-weight: 600;
        transition: all 0.2s ease;
        font-size: 14px;
    }
    .gr-button:hover {
        transform: translateY(-1px);
        box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);
    }
    .gr-form {
        background: white;
        border-radius: 12px;
        padding: 20px;
        margin: 10px 0;
    }
    h1, h2, h3 {
        color: #2c3e50;
        font-weight: 700;
    }
    /* Mobile Responsive */
    @media (max-width: 768px) {
        .gradio-container {
            padding: 10px 5px !important;
        }
        .main-container {
            padding: 5px;
        }
        .header-title {
            font-size: 1.8em !important;
        }
        .header-subtitle {
            font-size: 0.9em !important;
        }
        .feature-badges {
            flex-wrap: wrap !important;
            gap: 5px !important;
        }
        .feature-badge {
            font-size: 0.7em !important;
            padding: 3px 8px !important;
            margin: 2px !important;
        }
        .status-bar {
            flex-direction: column !important;
            gap: 10px !important;
            text-align: center !important;
        }
        .gr-button {
            font-size: 12px !important;
            padding: 8px 12px !important;
        }
    }
    @media (max-width: 480px) {
        .header-title {
            font-size: 1.5em !important;
        }
        .feature-badge {
            font-size: 0.6em !important;
        }
    }
    """
) as demo:

    # Header with improved mobile visibility
    gr.HTML("""
    <div class="header-section" style="margin-top: 0; padding-top: 0;">
        <div style="text-align: center; padding: 20px 15px; background: linear-gradient(135deg, #6c757d 0%, #495057 100%); color: white; border-radius: 15px; margin: 0 0 20px 0; box-shadow: 0 4px 12px rgba(0,0,0,0.15);">
            <h1 class="header-title" style="margin: 0; font-size: 2.2em; color: white; font-weight: 700; line-height: 1.2;">CFAx Agent</h1>
            <p class="header-subtitle" style="margin: 8px 0 0 0; font-size: 1.1em; opacity: 0.95; font-weight: 500;">Advanced Financial Analysis Platform</p>
            <p style="margin: 5px 0 0 0; font-size: 0.9em; opacity: 0.8;">AI-Powered • Real-Time Data • Code Generation</p>
            <div class="feature-badges" style="margin-top: 15px; display: flex; justify-content: center; align-items: center; flex-wrap: wrap; gap: 8px;">
                <span class="feature-badge" style="background: rgba(255,255,255,0.25); padding: 4px 12px; border-radius: 20px; font-size: 0.8em; font-weight: 500;">Live Data</span>
                <span class="feature-badge" style="background: rgba(255,255,255,0.25); padding: 4px 12px; border-radius: 20px; font-size: 0.8em; font-weight: 500;">AI Assistant</span>
                <span class="feature-badge" style="background: rgba(255,255,255,0.25); padding: 4px 12px; border-radius: 20px; font-size: 0.8em; font-weight: 500;">Code Execution</span>
                <span class="feature-badge" style="background: rgba(255,255,255,0.25); padding: 4px 12px; border-radius: 20px; font-size: 0.8em; font-weight: 500;">AI Visuals</span>
            </div>
        </div>
    </div>
    """)

    # Status indicator
    lib_status = sum([YFINANCE_AVAILABLE, SCIPY_AVAILABLE, SKLEARN_AVAILABLE, VISUALIZATION_AVAILABLE]) + 2
    openai_status = "✅ Connected" if openai_client else "❌ Not Connected"

    gr.HTML(f"""
    <div class="status-bar" style="display: flex; justify-content: space-between; align-items: center; background: white; padding: 12px 15px; border-radius: 10px; margin-bottom: 20px; box-shadow: 0 2px 8px rgba(0,0,0,0.1); font-size: 0.9em;">
        <div style="color: #000000;"><strong style="color: #000000;">Libraries:</strong> <span style="color: #28a745; font-weight: 600;">{lib_status}/6</span></div>
        <div style="color: #000000;"><strong style="color: #000000;">OpenAI:</strong> <span style="color: {'#28a745' if openai_client else '#dc3545'}; font-weight: 600;">{'✓' if openai_client else '✗'}</span></div>
        <div style="color: #000000;"><strong style="color: #000000;">Updated:</strong> <span style="color: #000000; font-weight: 500;">{datetime.now().strftime('%H:%M')}</span></div>
    </div>
    """)

    # Main tabs
    with gr.Tabs() as tabs:

        # AI Chat Assistant Tab
        with gr.Tab("AI Chat", id="chat"):
            chat_interface = gr.ChatInterface(
                fn=financial_chat,
                title="CFAx Financial Advisor",
                description="Ask me anything about financial analysis, market trends, or investment strategies!",
                examples=[
                    "Analyze AAPL stock performance",
                    "Explain Modern Portfolio Theory",
                    "Calculate DCF for a growth company",
                    "What are the best risk metrics?",
                    "Compare TSLA vs AAPL fundamentals",
                    "Show me portfolio optimization strategies"
                ],
                theme="soft"
            )

        # Stock Analysis Tab
        with gr.Tab("Stock Analysis", id="stocks"):
            gr.HTML("<h2 style='text-align: center; color: #2c3e50; font-size: 1.5em; margin: 10px 0;'>Real-Time Stock Analysis</h2>")

            with gr.Row():
                with gr.Column(scale=1):
                    stock_input = gr.Textbox(
                        label="Stock Symbol",
                        placeholder="Enter symbol (e.g., AAPL, TSLA, MSFT)",
                        value="AAPL"
                    )
                    period_select = gr.Dropdown(
                        choices=["1mo", "3mo", "6mo", "1y", "2y", "5y"],
                        value="1y",
                        label="Time Period"
                    )
                    analyze_btn = gr.Button(
                        "Analyze Stock",
                        variant="primary",
                        size="lg"
                    )

                with gr.Column(scale=2):
                    stock_chart = gr.Plot(label="Stock Chart")

            with gr.Row():
                stock_analysis = gr.Markdown(label="Analysis Results")

            def analyze_stock(symbol, period):
                if not symbol:
                    return None, "Please enter a stock symbol"

                stock_data = get_stock_data(symbol.upper(), period)
                if not stock_data or "error" in stock_data:
                    return None, f"Error fetching data for {symbol}"

                chart = create_stock_chart(symbol.upper(), stock_data["data"])
                analysis = perform_financial_analysis(symbol.upper())

                return chart, analysis

            analyze_btn.click(
                analyze_stock,
                inputs=[stock_input, period_select],
                outputs=[stock_chart, stock_analysis]
            )

        # Code Executor Tab
        with gr.Tab("Code Executor", id="code"):
            gr.HTML("<h2 style='text-align: center; color: #2c3e50; font-size: 1.5em; margin: 10px 0;'>Python Financial Code Executor</h2>")

            with gr.Row():
                with gr.Column():
                    code_input = gr.Code(
                        label="Python Code Editor",
                        language="python",
                        value="""# Financial Analysis Example
import numpy as np
import pandas as pd

# Calculate simple moving average
prices = [100, 102, 101, 103, 105, 104, 106]
ma_5 = sum(prices[-5:]) / 5
print(f"5-day Moving Average: ${ma_5:.2f}")

# Calculate returns
returns = [(prices[i] - prices[i-1]) / prices[i-1] for i in range(1, len(prices))]
print(f"Daily Returns: {[f'{r:.2%}' for r in returns]}")"""
                    )

                with gr.Column():
                    code_output = gr.Textbox(
                        label="Output",
                        lines=10,
                        max_lines=20
                    )

            with gr.Row():
                execute_btn = gr.Button(
                    "Execute Code",
                    variant="primary",
                    size="lg"
                )
                clear_btn = gr.Button("Clear", variant="secondary")

            def execute_code(code):
                if not code.strip():
                    return "Please enter some code to execute."
                return execute_python_code(code)

            def clear_code():
                return "", ""

            execute_btn.click(execute_code, inputs=[code_input], outputs=[code_output])
            clear_btn.click(clear_code, outputs=[code_input, code_output])

        # AI Visualizations Tab
        with gr.Tab("AI Visuals", id="visuals"):
            gr.HTML("<h2 style='text-align: center; color: #2c3e50; font-size: 1.5em; margin: 10px 0;'>AI-Generated Financial Visualizations</h2>")

            if openai_client:
                with gr.Row():
                    with gr.Column(scale=1):
                        viz_prompt = gr.Textbox(
                            label="Visualization Prompt",
                            placeholder="Describe the financial chart or visualization you want...",
                            lines=3
                        )
                        generate_btn = gr.Button(
                            "Generate Visualization",
                            variant="primary",
                            size="lg"
                        )

                    with gr.Column(scale=2):
                        viz_output = gr.Image(label="Generated Visualization")

                with gr.Row():
                    viz_examples = gr.Examples(
                        examples=[
                            ["Stock market bull vs bear visualization with charts"],
                            ["Modern portfolio dashboard with pie charts and graphs"],
                            ["Risk-return scatter plot with efficient frontier"],
                            ["Financial technology concept with AI and data"],
                            ["Cryptocurrency trading interface design"]
                        ],
                        inputs=[viz_prompt]
                    )

                def generate_viz(prompt):
                    if not prompt.strip():
                        return None

                    if not openai_client:
                        return None

                    image_url = generate_financial_visualization(prompt)
                    return image_url

                generate_btn.click(generate_viz, inputs=[viz_prompt], outputs=[viz_output])
            else:
                gr.HTML("""
                <div style="text-align: center; padding: 40px; background: #f8f9fa; border-radius: 10px;">
                    <h3 style="color: #6c757d;">OpenAI API Key Required</h3>
                    <p>Add your OpenAI API key to enable AI-generated visualizations</p>
                </div>
                """)

    # Footer
    gr.HTML("""
    <div style="text-align: center; padding: 15px; margin-top: 20px; background: #f8f9fa; border-radius: 10px; border-top: 3px solid #6c757d;">
        <p style="margin: 0; color: #6c757d; font-size: 0.9em;"><strong>Powered by:</strong> OpenAI GPT-4 • Python Libraries • Real-Time Data</p>
        <p style="margin: 5px 0 0 0; font-size: 0.8em; color: #adb5bd;">CFAx Agent - Advanced Financial Analysis Platform</p>
        <p style="margin: 8px 0 0 0; font-size: 0.7em; color: #adb5bd;">Designed + Developed by Lee Akpareva MBA, MA</p>
    </div>
    """)

if __name__ == "__main__":
    print("Starting CFAx Agent...")
    print(f"Libraries loaded: {sum([YFINANCE_AVAILABLE, SCIPY_AVAILABLE, SKLEARN_AVAILABLE, VISUALIZATION_AVAILABLE]) + 2}/6")
    demo.launch(share=True)