"""
eda_utils.py - Exploratory Data Analysis utilities for your stock app.
This file generates:
    - daily returns
    - returns histogram
    - drawdown curve
    - rolling volatility
    - correlation with NIFTY
    - basic indicators (SMA, RSI, MACD)
Everything is written simply for a beginner.
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import yfinance as yf

# ------------------------------------------------------------------
# BASIC RETURNS
# ------------------------------------------------------------------
def compute_daily_returns(df):
    """Return percent change (daily returns)."""
    returns = df['Close'].pct_change().dropna()
    return returns


# ------------------------------------------------------------------
# DRAWDOWN
# ------------------------------------------------------------------
def compute_drawdown(df):
    """Compute drawdown (drop from the peak)."""
    close = df['Close']
    rolling_max = close.cummax()
    drawdown = (close - rolling_max) / rolling_max
    return drawdown


# ------------------------------------------------------------------
# VOLATILITY
# ------------------------------------------------------------------
def compute_volatility(df, window=30):
    """Rolling 30-day standard deviation (volatility)."""
    return df['Close'].pct_change().rolling(window).std()


# ------------------------------------------------------------------
# CORRELATION WITH NIFTY
# ------------------------------------------------------------------
def correlation_with_nifty(df, start_date="2010-01-01"):
    """
    Downloads NIFTY (^NSEI) and computes correlation of returns.
    """
    nifty = yf.Ticker("^NSEI").history(start=start_date)
    if nifty.empty:
        return None

    stock_ret = df['Close'].pct_change()
    nifty_ret = nifty['Close'].pct_change()

    merged = pd.concat([stock_ret, nifty_ret], axis=1, join='inner').dropna()
    merged.columns = ['stock', 'nifty']

    return merged['stock'].corr(merged['nifty'])


# ------------------------------------------------------------------
# INDICATORS: SMA
# ------------------------------------------------------------------
def add_sma(df, windows=[20, 50, 200]):
    """Adds Simple Moving Averages to DataFrame."""
    for w in windows:
        df[f"SMA_{w}"] = df['Close'].rolling(w).mean()
    return df


# ------------------------------------------------------------------
# INDICATORS: RSI
# ------------------------------------------------------------------
def add_rsi(df, period=14):
    """Add RSI indicator."""
    delta = df['Close'].diff()
    gain = delta.clip(lower=0)
    loss = (-delta).clip(lower=0)

    avg_gain = gain.rolling(period).mean()
    avg_loss = loss.rolling(period).mean()

    rs = avg_gain / avg_loss
    df['RSI'] = 100 - (100 / (1 + rs))

    return df


# ------------------------------------------------------------------
# INDICATORS: MACD
# ------------------------------------------------------------------
def add_macd(df):
    """Add MACD, Signal Line."""
    exp1 = df['Close'].ewm(span=12, adjust=False).mean()
    exp2 = df['Close'].ewm(span=26, adjust=False).mean()
    df['MACD'] = exp1 - exp2
    df['Signal'] = df['MACD'].ewm(span=9, adjust=False).mean()
    return df


# ------------------------------------------------------------------
# EASY PLOTTING
# ------------------------------------------------------------------
def plot_eda(df, ticker):
    """Generate basic EDA plots."""

    plt.figure(figsize=(12,4))
    df['Close'].plot(title=f"{ticker} - Close Price")
    plt.grid(True)
    plt.show()

    plt.figure(figsize=(12,4))
    compute_daily_returns(df).hist(bins=40)
    plt.title("Histogram of Returns")
    plt.grid(True)
    plt.show()

    plt.figure(figsize=(12,4))
    compute_drawdown(df).plot(title="Drawdown Curve")
    plt.grid(True)
    plt.show()

    plt.figure(figsize=(12,4))
    compute_volatility(df).plot(title="Rolling Volatility (30-day)")
    plt.grid(True)
    plt.show()

    print("SMA added → columns:", [c for c in df.columns if 'SMA' in c])
    print("RSI added? →", 'RSI' in df.columns)
    print("MACD added? →", 'MACD' in df.columns)