# app.py
import ast
from datetime import datetime
import streamlit as st
import plotly.graph_objects as go
import pandas as pd
from utils.helper import *


# english
def main_algo_trader():
    # Front-end Design
    st.set_page_config(layout="wide")
    st.write("# Welcome to Algorithmic Trading! A Quick Implementation👋")

    with st.sidebar:
        with st.expander("Expand/Collapse"):
            st.markdown(
                r"""
                The following app is a simple demonstration of the growth stock strategy. For simplicity, we assume our research team hand over a pool of stocks. Amongst the pool of stocks, we can do the following:
                - use `yfinance` library to download stock data live (for the sake of speed, please start with time frame "1mo");
                - every period (time frame is a tuning parameter), we balance our portfolio (equal weight) by holding the top n stocks (n can be top quintile/quartile of stocks);
                """
            )

    # Main inputs
    tickers = st.text_input(
        "Enter tickers (comma-separated):",
        "MSFT, AAPL, NVDA, GOOG, AMZN, META, LLY, AVGO, TSLA, JPM, V, WMT, UNH, MA, PG, HD, JNJ, ORCL, MRK, COST, ABBV, BAC, CRM, AMD, NFLX, ACN, ADBE, DIS, TMO, WFC, MCD, CSCO, ABT, QCOM, INTC, INTU, IBM, AMAT, CMCSA, AXP, PFE, NOW, AMGN, MU",
    )
    start_date = st.sidebar.date_input("Start date", pd.to_datetime("2001-01-01"))
    end_date = st.sidebar.date_input(
        "End date", pd.to_datetime(datetime.now().strftime("%Y-%m-%d"))
    )
    time_frame = st.sidebar.selectbox(
        "Select Time Frame:",
        [
            "1mo",
            "3mo",
        ],
    )
    top_n = st.sidebar.number_input("Top n stocks", min_value=1, value=3)
    height_of_graph = st.sidebar.number_input(
        "Height of the plot", min_value=500, value=750
    )

    # Process inputs
    tickers_list = [ticker.strip() for ticker in tickers.split(",")]

    # Run analysis on button click
    if st.button("Run Analysis"):
        with st.spinner("Downloading data and calculating returns..."):
            stock_data = download_stock_data(
                tickers_list,
                start_date.strftime("%Y-%m-%d"),
                end_date.strftime("%Y-%m-%d"),
                w=time_frame,
            )
            returns_data = create_portfolio_and_calculate_returns(stock_data, top_n)
            benchmark_sharpe_ratio = (
                returns_data["benchmark"].mean() / returns_data["benchmark"].std()
            )
            portfolio_sharpe_ratio = (
                returns_data["portfolio_returns"].mean()
                / returns_data["portfolio_returns"].std()
            )

            # Data for plotting
            df = returns_data[
                ["rolling_benchmark", "rolling_portfolio_returns", "portfolio_history"]
            ]
            df.index = pd.to_datetime(df.index, unit="ms")

            # Create download file
            @st.cache_data
            def convert_df(df):
                # IMPORTANT: Cache the conversion to prevent computation on every rerun
                return df.to_csv().encode("utf-8")

            csv = convert_df(df)

            # Create plot
            fig = go.Figure()
            fig.add_trace(
                go.Scatter(
                    x=df.index,
                    y=df["rolling_benchmark"],
                    mode="lines",
                    name="Rolling Benchmark",
                )
            )
            fig.add_trace(
                go.Scatter(
                    x=df.index,
                    y=df["rolling_portfolio_returns"],
                    mode="lines",
                    name="Rolling Portfolio Returns",
                )
            )

            for date, stocks in df["portfolio_history"].items():
                fig.add_shape(
                    type="line",
                    x0=date,
                    y0=0,
                    x1=date,
                    y1=0,
                    line=dict(color="RoyalBlue", width=1),
                )
                fig.add_annotation(
                    x=date,
                    y=0.5,
                    text=str(stocks),
                    showarrow=False,
                    yshift=10,
                    textangle=-90,
                    font=dict(size=15),  # You can adjust the size as needed
                )

            # Calculate means and standard deviations
            benchmark_mean = returns_data["benchmark"].mean()
            benchmark_std = returns_data["benchmark"].std()
            portfolio_mean = returns_data["portfolio_returns"].mean()
            portfolio_std = returns_data["portfolio_returns"].std()

            # Update title text with additional information
            if time_frame == "1mo":
                some_n_based_on_time_frame = 12
                in_a_year = 1000 * (1 + portfolio_mean) ** (some_n_based_on_time_frame)
                in_50_years = 1000 * (1 + portfolio_mean) ** (
                    some_n_based_on_time_frame * 50
                )
            else:
                some_n_based_on_time_frame = 4
                in_a_year = 1000 * (1 + portfolio_mean) ** (some_n_based_on_time_frame)
                in_50_years = 1000 * (1 + portfolio_mean) ** (
                    some_n_based_on_time_frame * 50
                )
            title_text = (
                f"Performance:<br>"
                f"Benchmark Sharpe Ratio = {benchmark_sharpe_ratio:.3f}, "
                f"Portfolio Sharpe Ratio = {portfolio_sharpe_ratio:.3f}, "
                f"based on time frame: {time_frame}<br>"
                f"Benchmark => Mean: {benchmark_mean:.4f}, Std: {benchmark_std:.4f}; "
                f"Portfolio => Mean: {portfolio_mean:.4f}, Std: {portfolio_std:.4f}<br>"
                f"---<br>"
                f"This may or may not be a small number, let's check: <br>"
                f"$1,000*(1+{portfolio_mean:.4f})^({some_n_based_on_time_frame})={in_a_year}, <br>"
                f"$1,000*(1+{portfolio_mean:.4f})^({some_n_based_on_time_frame}*50)={in_50_years}."
            )
            curr_max_num = max(
                df.rolling_benchmark.max(), df.rolling_portfolio_returns.max()
            )
            fig.update_layout(
                title=title_text,
                xaxis_title="Date",
                yaxis_title="Value",
                yaxis=dict(range=[0, curr_max_num * 1.1]),
                legend=dict(
                    orientation="h", x=0.5, y=-0.4, xanchor="center", yanchor="bottom"
                ),
                height=height_of_graph,
            )

            st.plotly_chart(fig, use_container_width=True)

            # Post-analysis
            col1, col2 = st.columns(2)

            with col1:
                # Checkpoint: ask user whether they want portfolio weights
                if csv:
                    recent_selected_stocks = df["portfolio_history"][-1]
                    recent_selected_stocks = ", ".join(recent_selected_stocks)
                    st.success(
                        f"The algorithm suggests to hold the following stocks for the current month (equally weighted): {recent_selected_stocks}"
                    )

            with col2:
                # Download
                st.download_button(
                    label="Download data as CSV",
                    data=csv,
                    file_name=f"history_{end_date}.csv",
                    mime="text/csv",
                )

            # Question-answering
            if csv:
                user_question = st.text_input(
                    "Enter your question:",
                    "Tell me about the company with stock ticker AAPL.",
                )
                st.write(f"Question entered: {user_question}")
                with st.spinner("Running langchain agent with GPT3..."):
                    final_ans = run_langchain_agent_(
                        question=str(user_question),
                        interested_tickers=str(recent_selected_stocks),
                    )
                    st.markdown(final_ans)


# chinese
def main_algo_trader_chinese():
    # Front-end Design
    st.set_page_config(layout="wide")
    st.write("# 欢迎来到算法交易！一个快速模拟平台👋")

    with st.sidebar:
        with st.expander("Expand/Collapse"):
            st.markdown(
                r"""
                以下应用程序是成长股策略的简单演示。为简单起见，我们假设我们的研究团队交给了一些股票。在这些股票池中，我们可以做到以下几点：
                - 使用 `yfinance` 库实时下载股票数据（为了速度，请从时间框架 "1mo" 开始）；
                - 每个周期（时间框架是一个调整参数），我们通过持有前 n 只股票（n 可以是股票的前五分之一/四分之一）来平衡我们的投资组合（等权）。
                """
            )

    # Main inputs
    tickers = st.text_input(
        "使用英文键入输入股票代码（以逗号分隔）:",
        "MSFT, AAPL, NVDA, GOOG, AMZN, META, LLY, AVGO, TSLA, JPM, V, WMT, UNH, MA, PG, HD, JNJ, ORCL, MRK, COST, ABBV, BAC, CRM, AMD, NFLX, ACN, ADBE, DIS, TMO, WFC, MCD, CSCO, ABT, QCOM, INTC, INTU, IBM, AMAT, CMCSA, AXP, PFE, NOW, AMGN, MU",
    )
    start_date = st.sidebar.date_input("开始日期", pd.to_datetime("2001-01-01"))
    end_date = st.sidebar.date_input(
        "结束日期", pd.to_datetime(datetime.now().strftime("%Y-%m-%d"))
    )
    time_frame = st.sidebar.selectbox(
        "选择时间框架:",
        [
            "1mo",
            "3mo",
        ],
    )
    top_n = st.sidebar.number_input("选择前 n 支股票", min_value=1, value=3)
    height_of_graph = st.sidebar.number_input("图像高度", min_value=500, value=750)

    # Process inputs
    tickers_list = [ticker.strip() for ticker in tickers.split(",")]

    # Run analysis on button click
    if st.button("运行分析"):
        with st.spinner("下载数据并计算回报..."):
            stock_data = download_stock_data(
                tickers_list,
                start_date.strftime("%Y-%m-%d"),
                end_date.strftime("%Y-%m-%d"),
                w=time_frame,
            )
            returns_data = create_portfolio_and_calculate_returns(stock_data, top_n)
            benchmark_sharpe_ratio = (
                returns_data["benchmark"].mean() / returns_data["benchmark"].std()
            )
            portfolio_sharpe_ratio = (
                returns_data["portfolio_returns"].mean()
                / returns_data["portfolio_returns"].std()
            )

            # Data for plotting
            df = returns_data[
                ["rolling_benchmark", "rolling_portfolio_returns", "portfolio_history"]
            ]
            df.index = pd.to_datetime(df.index, unit="ms")

            # Create download file
            @st.cache_data
            def convert_df(df):
                # IMPORTANT: Cache the conversion to prevent computation on every rerun
                return df.to_csv().encode("utf-8")

            csv = convert_df(df)

            # Create plot
            fig = go.Figure()
            fig.add_trace(
                go.Scatter(
                    x=df.index,
                    y=df["rolling_benchmark"],
                    mode="lines",
                    name="Rolling Benchmark",
                )
            )
            fig.add_trace(
                go.Scatter(
                    x=df.index,
                    y=df["rolling_portfolio_returns"],
                    mode="lines",
                    name="Rolling Portfolio Returns",
                )
            )

            for date, stocks in df["portfolio_history"].items():
                fig.add_shape(
                    type="line",
                    x0=date,
                    y0=0,
                    x1=date,
                    y1=0,
                    line=dict(color="RoyalBlue", width=1),
                )
                fig.add_annotation(
                    x=date,
                    y=0.5,
                    text=str(stocks),
                    showarrow=False,
                    yshift=10,
                    textangle=-90,
                    font=dict(size=15),  # You can adjust the size as needed
                )

            # Calculate means and standard deviations
            benchmark_mean = returns_data["benchmark"].mean()
            benchmark_std = returns_data["benchmark"].std()
            portfolio_mean = returns_data["portfolio_returns"].mean()
            portfolio_std = returns_data["portfolio_returns"].std()

            # Update title text with additional information
            if time_frame == "1mo":
                some_n_based_on_time_frame = 12
                in_a_year = 1000 * (1 + portfolio_mean) ** (some_n_based_on_time_frame)
                in_50_years = 1000 * (1 + portfolio_mean) ** (
                    some_n_based_on_time_frame * 50
                )
            else:
                some_n_based_on_time_frame = 4
                in_a_year = 1000 * (1 + portfolio_mean) ** (some_n_based_on_time_frame)
                in_50_years = 1000 * (1 + portfolio_mean) ** (
                    some_n_based_on_time_frame * 50
                )
            title_text = (
                f"业绩:<br>"
                f"标杆回报风险比 = {benchmark_sharpe_ratio:.3f}, "
                f"投资组合回报风险比 = {portfolio_sharpe_ratio:.3f}, "
                f"交易窗口: {time_frame}<br>"
                f"标杆 => Mean: {benchmark_mean:.4f}, Std: {benchmark_std:.4f}; "
                f"投资组合 => Mean: {portfolio_mean:.4f}, Std: {portfolio_std:.4f}<br>"
                f"---<br>"
                f"这个数字如何理解，我们以1000块钱计算以下: <br>"
                f"$1,000*(1+{portfolio_mean:.4f})^({some_n_based_on_time_frame})={in_a_year}, <br>"
                f"$1,000*(1+{portfolio_mean:.4f})^({some_n_based_on_time_frame}*50)={in_50_years}."
            )
            curr_max_num = max(
                df.rolling_benchmark.max(), df.rolling_portfolio_returns.max()
            )
            fig.update_layout(
                title=title_text,
                xaxis_title="Date",
                yaxis_title="Value",
                yaxis=dict(range=[0, curr_max_num * 1.1]),
                legend=dict(
                    orientation="h", x=0.5, y=-0.4, xanchor="center", yanchor="bottom"
                ),
                height=height_of_graph,
            )

            st.plotly_chart(fig, use_container_width=True)

            # Post-analysis
            col1, col2 = st.columns(2)

            with col1:
                # Checkpoint: ask user whether they want portfolio weights
                if csv:
                    recent_selected_stocks = df["portfolio_history"][-1]
                    recent_selected_stocks = ", ".join(recent_selected_stocks)
                    st.success(f"算法建议在本月持有以下股票（均仓位）: {recent_selected_stocks}")

            with col2:
                # Download
                st.download_button(
                    label="下载数据为CSV格式",
                    data=csv,
                    file_name=f"history_{end_date}.csv",
                    mime="text/csv",
                )