import gradio as gr
import matplotlib.pyplot as plt
import pandas as pd
import sqlite3
from datetime import datetime
import sqlite3
from datetime import datetime

class Stock:
    def __init__(self, symbol, db_path, start_date, end_date):
        self.symbol = symbol
        self.db_path = db_path
        self.start_date = start_date
        self.end_date = end_date
        self.prices = self._download_prices()

    def _download_prices(self):
        conn = sqlite3.connect(self.db_path)
        query = f"SELECT Date, Close FROM `{self.symbol}` WHERE Date BETWEEN '{self.start_date}' AND '{self.end_date}'"
        prices = pd.read_sql_query(query, conn, parse_dates=['Date'])
        prices.set_index('Date', inplace=True)
        return prices

    def CurPrice(self, curDate):
        return self.prices.loc[curDate, 'Close'] if curDate in self.prices.index else None

    def NDayRet(self, N, curDate):
        if curDate not in self.prices.index:
            return None
        start_date = self.prices.index[self.prices.index.get_loc(curDate) - N]
        start_price = self.prices.loc[start_date, 'Close']
        end_price = self.prices.loc[curDate, 'Close']
        return (end_price - start_price) / start_price

    def DailyRet(self, curDate):
        if curDate not in self.prices.index:
            return None
        previous_date = self.prices.index[self.prices.index.get_loc(curDate) - 1]
        previous_price = self.prices.loc[previous_date, 'Close']
        current_price = self.prices.loc[curDate, 'Close']
        return (current_price - previous_price) / previous_price

    def Last30daysPrice(self, curDate):
        end_loc = self.prices.index.get_loc(curDate) + 1
        start_loc = max(0, end_loc - 30)
        return self.prices.iloc[start_loc:end_loc]['Close'].values




# Function to calculate monthly and daily percentage change
def calculate_percentage_changes(stocks):
    monthly_pct_change = pd.DataFrame()
    for symbol, stock_obj in stocks.items():
        monthly_pct_change[symbol] = stock_obj.prices['Close'].resample('M').last().pct_change()
    monthly_pct_change.fillna(0, inplace=True)
    return monthly_pct_change

# Function to update stock investments
def update_stock_investments(monthly_pct_change, month_index, stock_investments):
    month = monthly_pct_change.index[month_index]
    month_performance = monthly_pct_change.loc[month]
    positive_stocks = [stock for stock, pct_change in month_performance.items() if pct_change > 0]
    portfolio_value = sum(stock_investments.iloc[month_index - 1][symbol] * (1 + monthly_pct_change.at[month, symbol])
                          for symbol in nifty_50_symbols if pd.notna(monthly_pct_change.at[month, symbol]))

    if positive_stocks:
        investment_per_positive_stock = portfolio_value / len(positive_stocks)
        stock_investments.loc[month] = {stock: investment_per_positive_stock if stock in positive_stocks else 0 for stock in nifty_50_symbols}
    else:
        stock_investments.loc[month] = 0

    return portfolio_value


def calculate_portfolio(start_date, end_date, investment_amount):

  # Initialize stock dataframes and stocks object
  stocks = {symbol: Stock(symbol, db_path, start_date, end_date) for symbol in nifty_50_symbols}
  monthly_pct_change = calculate_percentage_changes(stocks)
  stock_investments = pd.DataFrame(index=monthly_pct_change.index, columns=nifty_50_symbols)
  portfolio_returns = pd.DataFrame(index=monthly_pct_change.index, columns=["portfolio_returns"])

  # Initialize stock investments
  num_stocks = len(nifty_50_symbols)
  investment_per_stock = investment_amount / num_stocks
  stock_investments.iloc[0] = investment_per_stock

  # Calculate portfolio returns
  for month_index in range(1, len(monthly_pct_change.index)):
      portfolio_value = update_stock_investments(monthly_pct_change, month_index, stock_investments)
      portfolio_returns.at[monthly_pct_change.index[month_index], 'portfolio_returns'] = portfolio_value 

  # Create Stock object and Calculate monthly returns for NIFTY50
  nifty_50_stock = Stock('NIFTY50', db_path, start_date, end_date)
  nifty_50_monthly_return = nifty_50_stock.prices['Close'].resample('M').last().pct_change()
  nifty_50_portfolio_change = nifty_50_monthly_return*(investment_amount)
  nifty_50_portfolio_value =nifty_50_portfolio_change.cumsum()+investment_amount

  return portfolio_returns, nifty_50_portfolio_value

def plot_chart(start_date, end_date, investment_amount):
  
  portfolio_returns, nifty_50_portfolio_value = calculate_portfolio(start_date, end_date, investment_amount)

  # Visualization
  plt.figure(figsize=(12, 6))
  portfolio_returns.plot()
  nifty_50_portfolio_value.plot()
  plt.title('Portfolio Returns Over Time')
  plt.xlabel('Date')
  plt.ylabel('Portfolio Value (in Rupees)')
  plt.legend()
  plt.grid(True)
  plt.savefig('portfolio_chart.png')
  plt.show()

  cagr = 100

  return 'portfolio_chart.png', f"CAGR: {cagr*100:.2f}%"

# Markdown text to be displayed
markdown_text = """
This interface allows you to test a stock selection strategy on nifty50 stocks and compare it with the return of nifty50 index. 
Enter the start date, end date, and the investment amount to get started.
Please pick a date range from 2000-01-01 to 2024-01-06
Please try shorter timeframes if you are facing an error """

# Gradio interface
iface = gr.Interface(
    fn=plot_chart,
    inputs=[
        gr.Textbox(label="Start Date (YYYY-MM-DD)"),
        gr.Textbox(label="End Date (YYYY-MM-DD)"),
        gr.Number(label="Investment Amount (in Rupees)")
    ],
    outputs=[
        gr.Image(type="filepath"),  # Corrected type for Image output
        gr.Textbox(label="CAGR")
    ],
    title="Stock Portfolio Performance",
    description=markdown_text  # Adding markdown text to the description
)

iface.launch()