import yfinance as yf
import pandas as pd

def process_dataframe(df):
    def get_rsi(close, lookback):
        ret = close.diff()
        up = []
        down = []
        for i in range(len(ret)):
            if ret[i] < 0:
                up.append(0)
                down.append(ret[i])
            else:
                up.append(ret[i])
                down.append(0)
        up_series = pd.Series(up)
        down_series = pd.Series(down).abs()
        up_ewm = up_series.ewm(com=lookback - 1, adjust=False).mean()
        down_ewm = down_series.ewm(com=lookback - 1, adjust=False).mean()
        rs = up_ewm / down_ewm
        rsi = 100 - (100 / (1 + rs))
        rsi_df = pd.DataFrame(rsi).rename(columns={0: 'RSI'}).set_index(close.index)
        rsi_df = rsi_df.dropna()
        return rsi_df[3:]

    df['RSI'] = get_rsi(df['Close'], 14)
    df['SMA20'] = df['Close'].rolling(window=20).mean()
    df.drop(['Adj Close'], axis=1, inplace=True)
    df = df.dropna()

    return df

def fin_data(ticker, startdate, enddate):
  df=yf.download(ticker,start=startdate,end=enddate, progress=False)
  df = process_dataframe(df)
  df.reset_index(inplace=True)
  df = df.dropna()
  df.reset_index(drop=True, inplace=True)
  df[['Open', 'High', 'Low', 'Close',"RSI"]] = df[['Open', 'High', 'Low', 'Close',"RSI"]].round(2)
  df = df[200:]
  df.reset_index(drop=True,inplace=True)
  return df

def eqt(ticker, startdate, enddate, share_qty = 90):
  df = fin_data(ticker, startdate, enddate)
  entry = False
  trading = False
  shares_held = 0
  buy_price = 0
  target1 = False
  target2 = False
  target3 = False
  tgt1 = 0
  tgt2 = 0
  tgt3 = 0
  total_profit = 0
  profits = []
  stop_loss = 0
  capital_list = []
  start_date = []
  end_date = []

  for i in range(1, len(df)-1):
      try:
          if df.at[i, 'RSI'] > 60 and df.at[i - 1, 'RSI'] < 60 and df.at[i, 'High'] < df.at[i + 1, 'High'] and not entry and not trading:
              buy_price = df.at[i, 'High']
              stop_loss = df.at[i, 'Low']
              start_date.append(df.at[i, 'Date'])
              capital = buy_price * share_qty
              capital_list.append(round(capital, 2))
              shares_held = share_qty
              entry = True
              trading = True

          if trading and not target1:
              if (df.at[i + 1, 'High'] - buy_price) >= 0.02 * buy_price:
                  stop_loss = buy_price
                  target1 = True
                  tgt1 = 0.02 * buy_price * (share_qty / 3)
                  shares_held -= (share_qty / 3)
                  total_profit = tgt1

          if trading and target1 and not target2:
              if (df.at[i + 1, 'High'] - buy_price) >= 0.04 * buy_price:
                  target2 = True
                  tgt2 = 0.04 * buy_price * (share_qty / 3)
                  total_profit += tgt2
                  shares_held -= (share_qty / 3)

          if trading and target2 and not target3:
              if (df.at[i + 1, 'Open'] < df.at[i + 1, 'SMA20'] < df.at[i + 1, 'Close']) or (df.at[i + 1, 'Open'] > df.at[i + 1, 'SMA20'] > df.at[i + 1, 'Close']):
                  stop_loss = df.at[i + 1, 'Low']
                  if df.at[i + 2, 'Low'] < stop_loss:
                      target3 = True
                      tgt3 = stop_loss * (share_qty / 3)
                      shares_held -= (share_qty / 3)
                      total_profit += tgt3

          if (df.at[i + 1, 'Low'] < stop_loss and trading):
              profit_loss = (shares_held * stop_loss) - (shares_held * buy_price)
              total_profit += profit_loss
              profits.append(total_profit)
              end_date.append(df.at[i, 'Date'])
              shares_held = 0
              buy_price = 0
              entry = False
              trading = False
              target1 = target2 = target3 = False
              tgt1 = tgt2 = tgt3 = 0
              total_profit = 0
              
        


      except IndexError:
          continue

  print("\n")
  print(f"Stock: {ticker} - From {df.at[1, 'Date']} to {df.at[len(df) - 1, 'Date']}")
  print(f"Required capital Range equity per trade: {round(capital_list[0],2)} ₹ - {round(capital_list[-1],2)} ₹")
  print("Duration Total Trading Profit:", round(sum(profits), 2),"₹")
  if profits:
    if len(start_date) > len(end_date):
        rr = len(end_date)
        df = pd.DataFrame({"Start" : start_date[:rr], "End": end_date, "profit" : profits, "Capital" : capital_list[:rr]})
        df['percentage'] = (df['profit'] / df['Capital']) * 100
        df['percentage'] = df['percentage'].apply(lambda x: f"{x:.2f}%" if x >= 0 else f"-{-x:.2f}%")
    else:
        df = pd.DataFrame({"Start" : start_date, "End": end_date, "profit" : profits, "Capital" : capital_list})
        df['percentage'] = (df['profit'] / df['Capital']) * 100
        df['percentage'] = df['percentage'].apply(lambda x: f"{x:.2f}%" if x >= 0 else f"-{-x:.2f}%")
    return df
  else:
    return 0