import streamlit as st
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import random

# Streamlit app layout
st.title("Estimating statistics from two sets of historical data")

with st.form("data_input_form"):
    N = st.number_input("Enter the number of data points for each set:", step=1)

    # Button to generate random data sets
    if st.form_submit_button("Generate Random Data Sets"):
        # Generate two sets of N random data points each
        set1 = [random.randint(0, 20) for _ in range(N)]
        set2 = [random.randint(0, 20) for _ in range(N)]

if "set1" in globals():
    # Calculating mean of each set
    set1_mean = np.mean(set1)
    set2_mean = np.mean(set2)

    st.subheader("Expected Values")
    st.write(f"Mean of Set 1: {set1_mean:,.2f}")
    st.write(f"Mean of Set 2: {set2_mean:,.2f}")

    # Calculating variances
    set1_var = np.var(set1)
    set2_var = np.var(set2)

    st.subheader("Variances")
    st.write(f"Variance of Set 1: {set1_var:,.2f}")
    st.write(f"Variance of Set 2: {set2_var:,.2f}")

    # Calculating correlation coefficient
    corr_coefficient = np.corrcoef(set1, set2)[0, 1]

    st.subheader("Correlation Coefficient")
    st.write(f"Correlation Coefficient: {corr_coefficient:,.2f}")

    # Scatter plot
    plt.scatter(set1, set2)
    plt.title("Scatter Plot")
    plt.xlabel("Set 1")
    plt.ylabel("Set 2")

    st.pyplot(plt)
    plt.close()

    # Line plot
    plt.title("Line Plot of Data Sets")
    plt.plot(np.arange(len(set1)), set1, label="Set 1")
    plt.plot(np.arange(len(set2)), set2, label="Set 2")
    plt.ylim(-10, 40)
    plt.xlabel("Time")
    plt.ylabel("Returns")
    plt.legend()

    st.pyplot(plt)
    plt.close()

    # Calculating weights and portfolio returns
    w_set1 = (set1_mean - set2_mean + set2_var - corr_coefficient * np.sqrt(set1_var * set2_var)) / (
        set1_var + set2_var - 2 * corr_coefficient * np.sqrt(set1_var * set2_var)
    )
    w_set2 = 1 - w_set1

    portfolio_var = (w_set1**2)*set1_var + (w_set2**2)*set2_var + 2*w_set1*w_set2*corr_coefficient*np.sqrt(set1_var * set2_var)

    st.subheader("Weights and Portfolio Returns")
    st.write("Assuming Set 1 and Set 2 represent returns of portfolios")
    st.write(f"Weight of Set 1: {w_set1:,.2f}")
    st.write(f"Weight of Set 2: {w_set2:,.2f}")
    st.write(f"Expected Portfolio Return: {w_set1*set1_mean + w_set2*set2_mean:,.2f}")
    st.write(f"Portfolio Variance: {portfolio_var:,.2f}")

else:
    st.write("Please input the number of data points for each set.")