import streamlit as st
import yfinance as yf
import pandas as pd
import matplotlib.pyplot as plt
import imageio
from io import BytesIO
import tempfile
import math

st.set_page_config(layout="wide")
st.title("Market Cap Animation")

tickers_input = st.text_input("Tickers", "AAPL,MSFT,NVDA,GOOG,AMZN,META")
start_date_str = st.text_input("Start Date", "2021-01-01")
end_date_str = st.text_input("End Date", "2025-04-01")

def fetch_market_caps(tickers, start, end):
    data = yf.download(tickers, start=start, end=end)["Close"]
    if isinstance(data, pd.Series):
        data = data.to_frame()
    mcap = pd.DataFrame(index=data.index)
    failed = []
    for t in tickers:
        try:
            shares = yf.Ticker(t).info.get("sharesOutstanding", None)
            if not shares or shares <= 0:
                failed.append(t)
                continue
            mcap[t] = data[t] * shares / 1e6
        except:
            failed.append(t)
    mcap.dropna(how="any", inplace=True)
    return mcap.iloc[::5], failed

if st.button("Generate"):
    tickers = [t.strip().upper() for t in tickers_input.split(",")]
    start = pd.to_datetime(start_date_str)
    end = pd.to_datetime(end_date_str)
    mcap, failed = fetch_market_caps(tickers, start, end)

    if mcap.empty:
        st.error("No data.")
    else:
        frames = []
        fig, ax = plt.subplots(figsize=(12, 6))

        ordered = list(mcap.columns)
        colors = {t: plt.cm.tab10(i % 10) for i, t in enumerate(ordered)}

        # Grid layout dimensions
        n = len(ordered)
        cols = math.ceil(math.sqrt(n))
        rows = math.ceil(n / cols)

        # Fixed positions in a grid
        grid_positions = {}
        for i, t in enumerate(ordered):
            row = i // cols
            col = i % cols
            x = col + 1
            y = rows - row
            grid_positions[t] = (x, y)

        for date, row in mcap.iterrows():
            ax.clear()
            ax.axis("off")
            ax.set_xlim(0, cols + 2)
            ax.set_ylim(0, rows + 2)
            ax.set_title(str(date.date()), fontsize=14)

            for t in ordered:
                val = row[t]
                x, y = grid_positions[t]
                ax.text(x, y, f"{t}\n{val:,.0f}M", fontsize=10,
                        ha='center', va='center',
                        bbox=dict(boxstyle="circle,pad=0.6", fc=colors[t], ec="black"))

            buf = BytesIO()
            plt.savefig(buf, format="png", bbox_inches="tight")
            buf.seek(0)
            frames.append(imageio.v2.imread(buf))
            buf.close()

        with tempfile.NamedTemporaryFile(delete=False, suffix=".gif") as tmp:
            imageio.mimsave(tmp.name, frames, duration=2.5)
            st.image(frames[-1])
            st.download_button("Download GIF", open(tmp.name, "rb").read(),
                               file_name="market_cap.gif", mime="image/gif")