"""
Portfolio Volatility Analyzer - Main Streamlit Application

Features:
- OCR parsing of portfolio screenshots
- Editable portfolio JSON
- Financial calculations (weights, returns, covariance, variance, volatility)
- Beautiful LaTeX formula displays for all calculations
- Interactive sliders for portfolio rebalancing
- Real-time recalculation
"""

import streamlit as st
from PIL import Image
import json

# Import our modules
import ocr_parser
import portfolio_calculator
import formula_generator


# Page configuration
st.set_page_config(
    page_title="Portfolio Volatility Analyzer",
    page_icon="📊",
    layout="wide",
    initial_sidebar_state="expanded"
)


# Initialize session state
if 'portfolio_data' not in st.session_state:
    st.session_state.portfolio_data = None
if 'portfolio_validated' not in st.session_state:
    st.session_state.portfolio_validated = False
if 'metrics' not in st.session_state:
    st.session_state.metrics = None
if 'show_all_terms' not in st.session_state:
    st.session_state.show_all_terms = False


# Main title and description
st.title("📊 Portfolio Volatility Analyzer with OCR")
st.markdown("""
Analyze your investment portfolio risk using **modern portfolio theory**.

**Features:**
- 📸 Upload portfolio screenshot for automatic OCR parsing
- ✏️ Edit portfolio data as JSON
- 📈 Fetch historical price data automatically
- 🧮 Calculate portfolio volatility with detailed mathematical formulas
- 🎚️ Interactive sliders for real-time portfolio rebalancing
""")

st.divider()


# ========================================
# Section 1: Portfolio Input
# ========================================

st.header("1️⃣ Portfolio Input")

# Create two columns for upload and manual entry
col1, col2 = st.columns([1, 1])

with col1:
    st.subheader("📸 Upload Screenshots")
    uploaded_files = st.file_uploader(
        "Upload one or more portfolio screenshots (PNG, JPG, JPEG)",
        type=["png", "jpg", "jpeg"],
        help="Upload screenshots of your portfolio. Multiple screenshots will be combined automatically.",
        accept_multiple_files=True,
        key="portfolio_uploader"
    )

    if uploaded_files:
        st.info(f"📤 Processing {len(uploaded_files)} screenshot(s)...")

        all_portfolios = []
        all_texts = []

        # Process each uploaded file
        for idx, uploaded_file in enumerate(uploaded_files, 1):
            st.markdown(f"### Screenshot {idx}")

            # Display uploaded image
            image = Image.open(uploaded_file)
            st.image(image, caption=f"Screenshot {idx}: {uploaded_file.name}")

            # OCR processing
            with st.spinner(f"Extracting text from screenshot {idx}..."):
                text, error = ocr_parser.extract_text_from_image(image)

            if error:
                st.error(f"❌ Screenshot {idx}: {error}")
                continue

            all_texts.append((idx, text))

            # Parse portfolio
            portfolio = ocr_parser.parse_portfolio(text)

            if portfolio:
                st.success(f"✅ Screenshot {idx}: Found {len(portfolio)} tickers: {', '.join(portfolio.keys())}")
                st.json(portfolio)
                all_portfolios.append(portfolio)
            else:
                st.warning(f"⚠️ Screenshot {idx}: No valid tickers found")

        # Show all extracted texts in an expander
        if all_texts:
            with st.expander("📄 View All Extracted Text"):
                for idx, text in all_texts:
                    st.markdown(f"**Screenshot {idx}:**")
                    st.text_area(f"OCR Output {idx}", text, height=100, disabled=True, key=f"ocr_text_{idx}")

        # Merge all portfolios
        if all_portfolios:
            merged_portfolio = ocr_parser.merge_portfolios(all_portfolios)
            st.success(f"✅ **Combined Portfolio:** {len(merged_portfolio)} unique tickers")
            st.json(merged_portfolio)
            st.session_state.portfolio_data = merged_portfolio
        else:
            st.warning("⚠️ **No valid tickers found in any screenshot.**")
            st.info("""
            **Possible reasons:**
            - Tickers are not in uppercase (e.g., 'aapl' instead of 'AAPL')
            - Company names instead of ticker symbols (e.g., 'Apple Inc.' instead of 'AAPL')
            - Unusual formatting or layout
            - Poor image quality

            **Solution:** Please manually enter your portfolio in the JSON editor below.
            """)
            st.session_state.portfolio_data = {}

with col2:
    st.subheader("✏️ Edit Portfolio (JSON)")

    st.info("""
    **Format:** `{"TICKER": amount, ...}`

    **Important:**
    - Use **ticker symbols** (e.g., AAPL, GOOGL, MSFT)
    - NOT company names (e.g., ❌ "Apple Inc.")
    - Tickers must be UPPERCASE
    - Amounts in your portfolio currency
    """)

    # Get initial JSON value
    if st.session_state.portfolio_data is not None and len(st.session_state.portfolio_data) > 0:
        initial_json = ocr_parser.format_portfolio_json(st.session_state.portfolio_data)
    else:
        # Default example
        initial_json = json.dumps({
            "AAPL": 5000,
            "GOOGL": 3000,
            "MSFT": 2000
        }, indent=2)

    # Editable text area
    edited_json = st.text_area(
        "Portfolio (JSON format)",
        value=initial_json,
        height=250,
        help="Edit the portfolio in JSON format: {\"TICKER\": amount, ...}"
    )

    # Validate button
    if st.button("✅ Validate Portfolio", type="primary"):
        is_valid, portfolio, error = ocr_parser.validate_portfolio_json(edited_json)

        if is_valid:
            st.session_state.portfolio_data = portfolio
            st.session_state.portfolio_validated = True
            st.success(f"✅ Portfolio validated! {len(portfolio)} tickers ready for analysis.")
        else:
            st.error(f"❌ {error}")
            st.session_state.portfolio_validated = False

st.divider()


# ========================================
# Section 2: Portfolio Analysis
# ========================================

if st.session_state.portfolio_validated and st.session_state.portfolio_data:

    st.header("2️⃣ Portfolio Analysis")

    portfolio = st.session_state.portfolio_data
    tickers = list(portfolio.keys())

    # Display current portfolio
    st.subheader("Current Portfolio")
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Tickers", len(tickers))
    with col2:
        total_value = sum(portfolio.values())
        st.metric("Total Value", f"${total_value:,.2f}")
    with col3:
        st.metric("Data Period", "1 year")

    # Fetch data and calculate metrics
    with st.spinner("🔄 Fetching historical data and calculating metrics..."):
        metrics, error = portfolio_calculator.get_portfolio_metrics(portfolio, period="1y")

    if error:
        st.error(f"❌ {error}")
        st.stop()

    # Store metrics in session state
    st.session_state.metrics = metrics

    st.success("✅ Analysis complete!")

    st.divider()

    # ========================================
    # Section 3: Data Display
    # ========================================

    st.header("3️⃣ Historical Data")

    # Portfolio Weights
    st.subheader("📊 Portfolio Weights")
    weights_df = [(ticker, f"{weight*100:.2f}%") for ticker, weight in metrics['weights'].items()]
    st.table(weights_df)

    # Historical Prices
    st.subheader("📈 Historical Prices (Last 5 Days)")
    st.dataframe(metrics['prices'].tail())

    # Returns
    with st.expander("📉 Daily Log Returns (Last 5 Days)"):
        st.dataframe(metrics['returns'].tail())

    # Covariance Matrix
    st.subheader("🔢 Covariance Matrix (Annualized)")
    st.dataframe(metrics['cov_matrix'] * 252)

    st.divider()

    # ========================================
    # Section 4: Mathematical Formulas
    # ========================================

    st.header("4️⃣ Mathematical Formulas")

    # Generate all formulas
    formulas = formula_generator.generate_all_formulas(
        amounts=portfolio,
        weights=metrics['weights'],
        cov_matrix=metrics['cov_matrix'],
        variance=metrics['variance'],
        volatility=metrics['volatility'],
        variance_breakdown=metrics['variance_breakdown']
    )

    # Weight Formulas
    st.subheader("⚖️ Portfolio Weights")
    st.markdown("**Symbolic Formula:**")
    st.latex(formulas['weights_symbolic'])
    st.markdown("**Numerical Calculation:**")
    st.latex(formulas['weights_numerical'])

    # Covariance Matrix
    st.subheader("📊 Covariance Matrix (Annualized)")
    st.latex(formulas['covariance_matrix'])

    # Correlation Matrix
    with st.expander("🔗 Correlation Matrix"):
        st.latex(formulas['correlation_matrix'])

    # Variance Formula
    st.subheader("📐 Portfolio Variance")
    st.markdown("**Symbolic Formula:**")
    st.latex(formulas['variance_symbolic'])

    st.markdown("**Detailed Expansion:**")
    st.latex(formulas['variance_expanded'])

    # Toggle for full expansion
    if st.checkbox("🔍 Show all variance terms (no truncation)", value=False):
        st.markdown("**Complete Expansion (All Terms):**")
        st.latex(formulas['variance_expanded_full'])

    # Volatility Formula
    st.subheader("📊 Portfolio Volatility")
    st.markdown("**Symbolic Formula:**")
    st.latex(formulas['volatility_symbolic'])
    st.markdown("**Numerical Result:**")
    st.latex(formulas['volatility_numerical'])

    st.divider()

    # ========================================
    # Section 5: Final Results
    # ========================================

    st.header("5️⃣ Final Results")

    col1, col2, col3 = st.columns(3)

    with col1:
        st.metric(
            label="Portfolio Variance",
            value=f"{metrics['variance']:.6f}",
            help="Annualized portfolio variance"
        )

    with col2:
        st.metric(
            label="Portfolio Volatility",
            value=f"{metrics['volatility']:.4f}",
            help="Annualized portfolio standard deviation (σ)"
        )

    with col3:
        st.metric(
            label="Volatility (%)",
            value=f"{metrics['volatility']*100:.2f}%",
            help="Annualized volatility as percentage"
        )

    st.divider()

    # ========================================
    # Section 6: Interactive Rebalancing
    # ========================================

    st.header("6️⃣ Interactive Portfolio Rebalancing")

    st.markdown("""
    **Adjust portfolio amounts** using the sliders below to see how volatility changes in real-time.
    """)

    # Create sliders for each ticker
    new_amounts = {}
    slider_cols = st.columns(min(len(tickers), 3))  # Max 3 columns

    for idx, ticker in enumerate(tickers):
        col_idx = idx % len(slider_cols)
        with slider_cols[col_idx]:
            original_amount = portfolio[ticker]
            new_amount = st.slider(
                f"{ticker}",
                min_value=0.0,
                max_value=original_amount * 3,  # Allow up to 3x original
                value=original_amount,
                step=100.0,
                format="$%.0f",
                key=f"slider_{ticker}"
            )
            new_amounts[ticker] = new_amount

    # Check if amounts changed
    amounts_changed = any(new_amounts[t] != portfolio[t] for t in tickers)

    if amounts_changed:
        st.subheader("🔄 Recalculated Metrics")

        # Recalculate with new amounts
        with st.spinner("Recalculating..."):
            new_metrics, error = portfolio_calculator.get_portfolio_metrics(new_amounts, period="1y")

        if error:
            st.error(f"❌ {error}")
        else:
            # Display new results
            col1, col2 = st.columns(2)

            with col1:
                st.markdown("**New Portfolio Weights:**")
                for ticker, weight in new_metrics['weights'].items():
                    st.write(f"{ticker}: {weight*100:.2f}%")

            with col2:
                st.markdown("**New Volatility:**")
                st.metric(
                    label="Updated Volatility",
                    value=f"{new_metrics['volatility']*100:.2f}%",
                    delta=f"{(new_metrics['volatility'] - metrics['volatility'])*100:.2f}%",
                    delta_color="inverse"  # Lower volatility is better
                )

else:
    # Show instructions if portfolio not validated
    st.info("👆 Please upload a portfolio screenshot or enter portfolio data above, then click 'Validate Portfolio' to begin analysis.")

st.divider()

# ========================================
# Footer
# ========================================

st.markdown("---")
st.markdown("""
<div style='text-align: center; color: gray;'>
    <p>Built with ❤️ using Streamlit | Powered by Modern Portfolio Theory</p>
    <p><small>Data source: Yahoo Finance (yfinance) | OCR: Tesseract</small></p>
</div>
""", unsafe_allow_html=True)