import gradio as gr
import yfinance as yf
import numpy as np
import pandas as pd
import plotly.graph_objects as go
from datetime import datetime
from io import BytesIO
import base64

def download_stock_data(ticker, start_date, end_date):
    stock = yf.Ticker(ticker)
    df = stock.history(start=start_date, end=end_date)
    return df

def plot_interactive_chart(ticker, start_date, end_date, chart_type):
    try:
        stock = yf.Ticker(ticker)
        data = stock.history(start=start_date, end=end_date)

        if data.empty:
            return "No data available for the specified date range."

        if chart_type == "Logarithmic":
            fig = create_logarithmic_chart(data, ticker)
        else:
            fig = create_candlestick_chart(data, ticker)

        return fig
    except Exception as e:
        return f"An error occurred: {str(e)}"

def create_logarithmic_chart(data, ticker):
    x = (data.index - data.index[0]).days
    y = np.log(data['Close'])
    slope, intercept = np.polyfit(x, y, 1)

    future_days = 365 * 10
    all_days = np.arange(len(x) + future_days)
    log_trend = np.exp(intercept + slope * all_days)

    inner_upper_band = log_trend * 2
    inner_lower_band = log_trend / 2
    outer_upper_band = log_trend * 4
    outer_lower_band = log_trend / 4

    extended_dates = pd.date_range(start=data.index[0], periods=len(all_days), freq='D')

    fig = go.Figure()

    fig.add_trace(go.Scatter(x=data.index, y=data['Close'], mode='lines', name='Close Price', line=dict(color='blue')))
    fig.add_trace(go.Scatter(x=extended_dates, y=log_trend, mode='lines', name='Log Trend', line=dict(color='red')))
    fig.add_trace(go.Scatter(x=extended_dates, y=inner_upper_band, mode='lines', name='Inner Upper Band', line=dict(color='green')))
    fig.add_trace(go.Scatter(x=extended_dates, y=inner_lower_band, mode='lines', name='Inner Lower Band', line=dict(color='green')))
    fig.add_trace(go.Scatter(x=extended_dates, y=outer_upper_band, mode='lines', name='Outer Upper Band', line=dict(color='orange')))
    fig.add_trace(go.Scatter(x=extended_dates, y=outer_lower_band, mode='lines', name='Outer Lower Band', line=dict(color='orange')))

    fig.update_layout(
        title=f'{ticker} Stock Price (Logarithmic Scale) with Extended Trend Lines and Outer Bands',
        xaxis_title='Date',
        yaxis_title='Price (Log Scale)',
        yaxis_type="log",
        height=800,
        legend=dict(x=0.01, y=0.99, bgcolor='rgba(255, 255, 255, 0.8)'),
        hovermode='x unified'
    )

    fig.update_xaxes(
        rangeslider_visible=True,
        rangeselector=dict(
            buttons=list([
                dict(count=1, label="1m", step="month", stepmode="backward"),
                dict(count=6, label="6m", step="month", stepmode="backward"),
                dict(count=1, label="YTD", step="year", stepmode="todate"),
                dict(count=1, label="1y", step="year", stepmode="backward"),
                dict(step="all")
            ])
        )
    )

    return fig

def create_candlestick_chart(data, ticker):
    fig = go.Figure(data=[go.Candlestick(x=data.index,
                open=data['Open'],
                high=data['High'],
                low=data['Low'],
                close=data['Close'])])

    fig.update_layout(
        title=f'{ticker} Stock Price (Candlestick Chart)',
        xaxis_title='Date',
        yaxis_title='Price',
        height=800,
        hovermode='x unified'
    )

    fig.update_xaxes(
        rangeslider_visible=True,
        rangeselector=dict(
            buttons=list([
                dict(count=1, label="1m", step="month", stepmode="backward"),
                dict(count=6, label="6m", step="month", stepmode="backward"),
                dict(count=1, label="YTD", step="year", stepmode="todate"),
                dict(count=1, label="1y", step="year", stepmode="backward"),
                dict(step="all")
            ])
        )
    )

    return fig

def export_data(ticker, start_date, end_date, format):
    data = download_stock_data(ticker, start_date, end_date)
    
    if format == 'CSV':
        output = BytesIO()
        data.to_csv(output, index=True)
        b64 = base64.b64encode(output.getvalue()).decode()
        return f'data:text/csv;base64,{b64}'
    elif format == 'Excel':
        output = BytesIO()
        with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
            data.to_excel(writer, sheet_name='Stock Data', index=True)
        b64 = base64.b64encode(output.getvalue()).decode()
        return f'data:application/vnd.openxmlformats-officedocument.spreadsheetml.sheet;base64,{b64}'
    elif format == 'PDF':
        output = BytesIO()
        fig = go.Figure(data=[go.Table(
            header=dict(values=list(data.columns)),
            cells=dict(values=[data[col] for col in data.columns])
        )])
        fig.write_image(output, format='pdf')
        b64 = base64.b64encode(output.getvalue()).decode()
        return f'data:application/pdf;base64,{b64}'

# Get the current date
current_date = datetime.now().strftime("%Y-%m-%d")

# Custom CSS to make charts full-width and larger, and apply additional styling
custom_css = """
.container {max-width: 100% !important; padding: 0 !important;}
.plot-container {height: 800px !important; width: 100% !important;}
.react-plotly-container {height: 100% !important; width: 100% !important;}

body {
    background-color: #f0f0f0;
    font-family: 'Helvetica', sans-serif;
}
.gradio-container {
    margin: auto;
    padding: 15px;
    border-radius: 10px;
    background-color: white;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}
.gr-button {
    background-color: #4CAF50;
    border: none;
    color: white;
    text-align: center;
    text-decoration: none;
    display: inline-block;
    font-size: 16px;
    margin: 4px 2px;
    cursor: pointer;
    border-radius: 5px;
    padding: 10px 24px;
}
.gr-button:hover {
    background-color: #45a049;
}
.gr-form {
    border: 1px solid #ddd;
    border-radius: 5px;
    padding: 15px;
    margin-bottom: 20px;
}
.gr-input {
    width: 100%;
    padding: 12px 20px;
    margin: 8px 0;
    display: inline-block;
    border: 1px solid #ccc;
    border-radius: 4px;
    box-sizing: border-box;
}
.gr-input:focus {
    border-color: #4CAF50;
    outline: none;
}
.gr-label {
    font-weight: bold;
    margin-bottom: 5px;
}
"""

# Create Gradio interface with updated theme
with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title="Log Charting Tool") as iface:
    gr.Markdown("# Log Charting Tool")
    gr.Markdown("Enter a stock ticker and date range to generate a chart and export data.")
    
    with gr.Row():
        ticker = gr.Textbox(label="Stock Ticker", value="MSFT")
        start_date = gr.Textbox(label="Start Date", value="2015-01-01")
        end_date = gr.Textbox(label="End Date", value=current_date)
    
    with gr.Row():
        chart_type = gr.Radio(["Logarithmic", "Candlestick"], label="Chart Type", value="Logarithmic")
        export_format = gr.Dropdown(["CSV", "Excel", "PDF"], label="Export Format", value="CSV")
    
    with gr.Row():
        generate_button = gr.Button("Generate Chart")
        export_button = gr.Button("Export Data")
    
    with gr.Row():
        chart_output = gr.Plot(label="Stock Chart")
        export_output = gr.File(label="Exported Data")
    
    generate_button.click(
        plot_interactive_chart,
        inputs=[ticker, start_date, end_date, chart_type],
        outputs=[chart_output]
    )
    
    export_button.click(
        export_data,
        inputs=[ticker, start_date, end_date, export_format],
        outputs=[export_output]
    )

# Launch the app
iface.launch()