import yfinance as yf
import matplotlib.pyplot as plt
import numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from PIL import Image
import io
import gradio as gr
from cachetools import cached, TTLCache
import cProfile
import pstats

# Global fontsize variable
FONT_SIZE = 32
# Company ticker mapping
COMPANY_TICKERS = { 
    'Union Pacific': 'UNP',
    'Canadian Pacific KC': 'CP',
    'FedEx': 'FDX',
    'Autozone': 'AZO',
    'XPO Logistics': 'XPO',
    'JB Hunt Transport': 'JBHT',
    'Old Dominion FL': 'ODFL',
    'Broadcom Inc':'AVGO',
    'Genuine Parts Co': 'GPC',
    'C.H. Robinson': 'CHRW',
    'Expeditors Int': 'EXPD', 
    'Landstar System': 'LSTR',
    'Saia': 'SAIA',
    'Knight-Swift Transportation': 'KNX',
    'Schneider National': 'SNDR',
    'Ryder System': 'R',
    'Tesla': 'TSLA',
    'Amazon': 'AMZN',
    'A.O. Smith': 'AOS',
    'Acushnet Holdings': 'GOLF',
    'Allison Transmission': 'ALSN',
    'AMETEK': 'AME',
    'AMN Healthcare': 'AMN',
    'Analog Devices': 'ADI',
    'Ansys': 'ANSS',
    'AptarGroup': 'ATR',
    'Aramark': 'ARMK',
    'Snap-On': 'SNA',
    'ArcBest': 'ARCB',
    'Arch Capital Group': 'ACGL',
    'Atlassian': 'TEAM',
    'AutoNation': 'AN',
    'Avnet': 'AVT',
    'Brookfield Renewable Partners': 'BEP',
    'Cadence Bank': 'CADE',
    'CACI International': 'CACI',
    'California Water Service': 'CWT',
    'Cambrex': 'CBM',
    'Capri Holdings': 'CPRI',
    'Carlisle Companies': 'CSL',
    'Catalent': 'CTLT',
    'CDK Global': 'CDK',
    'Celanese': 'CE',
    'Celsius Holdings': 'CELH',
    'Centene': 'CNC',
    'Central Garden & Pet': 'CENT',
    'Chart Industries': 'GTLS',
    'Chemed': 'CHE',
    'Cheniere Energy': 'LNG',
    'Chesapeake Energy': 'CHK',
    'Church & Dwight': 'CHD',
    'Cimarex Energy': 'XEC',
    'Cincinnati Financial': 'CINF',
    'Cinemark': 'CNK',
    'Cirrus Logic': 'CRUS',
    'Cloudflare': 'NET',
    'Coca-Cola Consolidated': 'COKE',
    'Comerica': 'CMA',
    'Commercial Metals': 'CMC',
    'CommScope': 'COMM',
    'Community Health Systems': 'CYH',
    'Compass Minerals': 'CMP',
    'Comstock Resources': 'CRK',
    'Conagra Brands': 'CAG',
    'Consolidated Communications': 'CNSL',
    'Cooper-Standard': 'CPS',
    'Copart': 'CPRT',
    'CoreLogic': 'CLGX',
    'Core-Mark': 'CORE',
    'Cousins Properties': 'CUZ',
    'Covenant Logistics': 'CVLG',
    'Cree': 'CREE',
    'Cullen/Frost Bankers': 'CFR',
    'Curtiss-Wright': 'CW',
    'CyrusOne': 'CONE',
    'D.R. Horton': 'DHI',
    'Daseke': 'DSKE',
    'Deckers Outdoor': 'DECK',
    'Del Taco Restaurants': 'TACO',
    'Deluxe': 'DLX',
    'Dentsply Sirona': 'XRAY',
    'Dorman Products': 'DORM',
    'Douglas Emmett': 'DEI',
    'Dover': 'DOV',
    'DuPont de Nemours': 'DD',
    'Dycom Industries': 'DY',
    'Eagle Materials': 'EXP',
    'East West Bancorp': 'EWBC',
    'Eaton Vance': 'EV',
    'Echo Global Logistics': 'ECHO',
    'Ecolab': 'ECL',
    'Edgewell Personal Care': 'EPC',
    'eHealth': 'EHTH',
    'Elanco Animal Health': 'ELAN',
    'Elbit Systems': 'ESLT',
    'EMCOR Group': 'EME',
    'Encompass Health': 'EHC',
    'Encore Capital Group': 'ECPG',
    'Endo International': 'ENDP',
    'Entegris': 'ENTG',
    'Envestnet': 'ENV',
    'EPAM Systems': 'EPAM',
    'EPR Properties': 'EPR',
    'EQT': 'EQT',
    'Equitrans Midstream': 'ETRN',
    'Everbridge': 'EVBG',
    'Evergy': 'EVRG',
    'Eversource Energy': 'ES',
    'Exelixis': 'EXEL',
    'Exponent': 'EXPO',
    'Express': 'EXPR',
    'Exterran': 'EXTN',
    'Exxon Mobil': 'XOM',
    'FactSet': 'FDS',
    'Fair Isaac': 'FICO',
    'Federal Realty': 'FRT',
    'Federated Hermes': 'FHI',
    'Ferro': 'FOE',
    'First American': 'FAF',
    'Fortune Brands Home & Security': 'FBHS',
    'Franklin Electric': 'FELE',
    'Fresenius Medical Care': 'FMS',
    'Fresh Del Monte Produce': 'FDP',
    'Fulton Financial': 'FULT',
    'Gartner': 'IT',
    'Genpact': 'G',
    'Gibraltar Industries': 'ROCK',
    'Gilead Sciences': 'GILD',
    'Glacier Bancorp': 'GBCI',
    'Global Payments': 'GPN',
    'Globant': 'GLOB',
    'Graphic Packaging Holding': 'GPK',
    'HD Supply': 'HDS',
    'Heico': 'HEI',
    'Helmerich & Payne': 'HP',
    'Henry Schein': 'HSIC',
    'Hess': 'HES',
    'Oracle': 'ORCL',
    'Uber': 'UBER',
    'Werner Enterprises': 'WERN'
}


# Cache with 1-day TTL
cache = TTLCache(maxsize=100, ttl=86400)

@cached(cache)
def fetch_historical_data(ticker, start_date, end_date):
    """Fetch historical stock data and market cap from Yahoo Finance."""
    try:
        data = yf.download(ticker, start=start_date, end=end_date)
        if data.empty:
            raise ValueError(f"No data found for ticker {ticker}")
        info = yf.Ticker(ticker).info
        market_cap = info.get('marketCap', 'N/A')
        if market_cap != 'N/A':
            market_cap = market_cap / 1e9  # Convert to billions
        return data, market_cap
    except Exception as e:
        print(f"Error fetching data for {ticker}: {e}")
        return None, 'N/A'

def plot_to_image(plt, title, market_cap):
    """Convert plot to a PIL Image object."""
    plt.title(title, fontsize=FONT_SIZE + 1, pad=40)
    plt.suptitle(f'Market Cap: ${market_cap:.2f} Billion', fontsize=FONT_SIZE - 5, y=0.92, weight='bold')
    plt.legend(fontsize=FONT_SIZE)
    plt.xlabel('Date', fontsize=FONT_SIZE)
    plt.ylabel('', fontsize=FONT_SIZE)
    plt.grid(True)
    plt.xticks(rotation=45, ha='right', fontsize=FONT_SIZE)
    plt.yticks(fontsize=FONT_SIZE)
    plt.tight_layout(rect=[0, 0, 1, 0.88])

    buf = io.BytesIO()
    plt.savefig(buf, format='png', dpi=400)
    plt.close()
    buf.seek(0)
    return Image.open(buf)

def plot_indicator(data, company_name, ticker, indicator, market_cap):
    """Plot selected technical indicator for a single company."""
    plt.figure(figsize=(16, 10))
    if indicator == "SMA":
        sma_55 = data['Close'].rolling(window=55).mean()
        sma_100 = data['Close'].rolling(window=100).mean()  # 100-day SMA
        sma_200 = data['Close'].rolling(window=252).mean()
        plt.plot(data.index, data['Close'], label='Close')
        plt.plot(data.index, sma_55, label='55-day SMA')
        plt.plot(data.index, sma_100, label='100-day SMA')  # Plot 100-day SMA
        plt.plot(data.index, sma_200, label='252-day SMA')
        plt.ylabel('Price', fontsize=FONT_SIZE)
    elif indicator == "MACD":
        exp1 = data['Close'].ewm(span=12, adjust=False).mean()
        exp2 = data['Close'].ewm(span=26, adjust=False).mean()
        macd = exp1 - exp2
        signal = macd.ewm(span=9, adjust=False).mean()
        plt.plot(data.index, macd, label='MACD')
        plt.plot(data.index, signal, label='Signal Line')
        plt.bar(data.index, macd - signal, label='MACD Histogram')
        plt.ylabel('MACD', fontsize=FONT_SIZE)

    return plot_to_image(plt, f'{company_name} ({ticker}) {indicator}', market_cap)

def plot_indicators(company_names, indicator_types):
    """Plot the selected indicators for the selected companies."""
    images = []
    total_market_cap = 0
    if len(company_names) > 7:
        return None, "You can select up to 7 companies at the same time.", None
    if len(company_names) > 1 and len(indicator_types) > 1:
        return None, "You can only select one indicator when selecting multiple companies.", None

    with ThreadPoolExecutor() as executor:
        future_to_company = {
            executor.submit(fetch_historical_data, COMPANY_TICKERS[company], '2000-01-01', datetime.now().strftime('%Y-%m-%d')): (company, indicator)
            for company in company_names
            for indicator in indicator_types
        }

        for future in as_completed(future_to_company):
            company, indicator = future_to_company[future]
            ticker = COMPANY_TICKERS[company]
            data, market_cap = future.result()
            if data is None:
                continue
            images.append(plot_indicator(data, company, ticker, indicator, market_cap))
            if market_cap != 'N/A':
                total_market_cap += market_cap

    return images, "", total_market_cap

def select_all_indicators(select_all):
    """Select or deselect all indicators based on the select_all flag."""
    indicators = ["SMA", "MACD"]
    return indicators if select_all else []

def launch_gradio_app():
    """Launch the Gradio app for interactive plotting."""
    company_choices = list(COMPANY_TICKERS.keys())
    indicators = ["SMA", "MACD"]

    def fetch_and_plot(company_names, indicator_types):
        images, error_message, total_market_cap = plot_indicators(company_names, indicator_types)
        if error_message:
            return [None] * len(indicator_types), error_message, None
        return images, "", f"Total Market Cap: ${total_market_cap:.2f} Billion" if total_market_cap else "N/A"

    with gr.Blocks() as demo:
        company_checkboxgroup = gr.CheckboxGroup(choices=company_choices, label="Select Companies")
        
        select_all_checkbox = gr.Checkbox(label="Select All Indicators", value=False, interactive=True)
        indicator_types_checkboxgroup = gr.CheckboxGroup(choices=indicators, label="Select Technical Indicators")
        select_all_checkbox.change(select_all_indicators, inputs=select_all_checkbox, outputs=indicator_types_checkboxgroup)
        
        plot_gallery = gr.Gallery(label="Indicator Plots")
        error_markdown = gr.Markdown()
        market_cap_text = gr.Markdown()

        gr.Interface(
            fetch_and_plot, 
            [company_checkboxgroup, indicator_types_checkboxgroup], 
            [plot_gallery, error_markdown, market_cap_text]
        )

    demo.launch()

def profile_code():
    """Profile the main functions to find speed bottlenecks."""
    profiler = cProfile.Profile()
    profiler.enable()

    launch_gradio_app()

    profiler.disable()
    stats = pstats.Stats(profiler).sort_stats('cumtime')
    stats.print_stats(10)

if __name__ == "__main__":
    profile_code()