#!/usr/bin/env python3
import logging
from typing import List, Tuple, Optional # For type hinting

import matplotlib.pyplot as plt # Base plotting library
import mplfinance as mpf        # Financial charting
import numpy as np             # Numerical operations (arrays)
import pandas as pd            # Data manipulation (DataFrames)
import yfinance as yf          # Downloading market data
# --- Setup logging ---
logging.basicConfig(
    level=logging.INFO, # Set minimum level of messages to show
    format='[%(asctime)s] %(levelname)s: %(message)s' # Message format
)

if __name__ == '__main__':
    # --- Configuration for the run ---
    TICKER = 'BTC-USD'  # Example: Bitcoin USD
    PERIOD = '1mo'     # Example: Data for the last 1 month
                         # Other options: '1y', '6mo', '5d', 'max', etc.
    SIGMA = 0.03       # Example: 3% reversal threshold for Zig Zag

# --- Download Data ---
    logging.info(f"Downloading {TICKER} data for period: {PERIOD}")
    try:
        # Download data, progress=False hides the download status bar
        data = yf.download(TICKER, period=PERIOD, progress=False)
        if data.empty:
            raise ValueError(f"No data downloaded for {TICKER}.")
        # yfinance sometimes returns multi-level columns, flatten them
        # data.columns = data.columns.droplevel(1) # Adjusted to be optional/checked
        if isinstance(data.columns, pd.MultiIndex):
             data.columns = data.columns.droplevel(1)
        logging.info(f"Data downloaded successfully. Shape: {data.shape}")
    except Exception as e:
        logging.exception(f"Failed to download data for {TICKER}: {e}")
        exit() # Stop script if data download fails
    # ... (Calculation and Plotting calls follow) ...

    def find_zigzag_extremes(
        high: np.ndarray,
        low: np.ndarray,
        close: np.ndarray,
        sigma: float
    ) -> pd.DataFrame:
        """ Docstring explaining function... """
        if len(close) < 2: # Need at least two points
             return pd.DataFrame(columns=['price', 'type', 'confirmation_idx'])
        # ... rest of the function ...

    # Inside find_zigzag_extremes:
    # State variables initialization...
    last_extreme_type = 0
    up_trend = None
    # ... other variables ...
    
    # Initialize based on first two points trend
    if close[1] > close[0]:
        up_trend = True
        last_extreme_type = -1 # Implied previous bottom
        candidate_extreme_price = high[0]
        candidate_extreme_idx = 0
        last_confirmed_extreme_price = low[0] # Record initial point data
        last_confirmed_extreme_idx = 0
    else:
        up_trend = False
        last_extreme_type = 1 # Implied previous top
        candidate_extreme_price = low[0]
        candidate_extreme_idx = 0
        last_confirmed_extreme_price = high[0] # Record initial point data
        last_confirmed_extreme_idx = 0

    # Inside the loop `for i in range(1, len(close)):`
    
    # Example: Checking for Top Confirmation
    if up_trend: # Looking for a top
        if high[i] >= candidate_extreme_price:
            # New higher high found, update candidate top
            candidate_extreme_price = high[i]
            candidate_extreme_idx = i
        elif close[i] < candidate_extreme_price * (1 - sigma): # Price reversed > sigma%
            # Price reversed significantly, confirm the candidate top
            if last_extreme_type != 1: # Avoid confirming consecutive tops
                logging.debug(f"Idx {i}: Top confirmed: Idx={candidate_extreme_idx}, Price={candidate_extreme_price:.2f}...")
                extremes_data.append((candidate_extreme_idx, candidate_extreme_price, 1, i)) # Store confirmed top (type 1)
                last_extreme_type = 1
                last_confirmed_extreme_price = candidate_extreme_price
                last_confirmed_extreme_idx = candidate_extreme_idx
                # Start looking for a bottom
                up_trend = False
                candidate_extreme_price = low[i] # Initial candidate bottom
                candidate_extreme_idx = i
    else: # Looking for a bottom
        # Similar logic for finding lower lows and checking for upward reversal:
        # if low[i] <= candidate_extreme_price: ...
        # elif close[i] > candidate_extreme_price * (1 + sigma): ...
        #     if last_extreme_type != -1: ... # Confirm bottom (type -1)
            # ... update state, switch to looking for top ...

    # After the loop in find_zigzag_extremes:
    if not extremes_data:
        logging.warning("No significant extremes found...")
        return pd.DataFrame(columns=['price', 'type', 'confirmation_idx'])
    
    # Create DataFrame
    extremes_df = pd.DataFrame(
        extremes_data,
        columns=['extreme_idx', 'price', 'type', 'confirmation_idx']
    )
    extremes_df = extremes_df.set_index('extreme_idx') # Use bar number as index
    extremes_df = extremes_df.sort_index()
    # Remove duplicates if any (unlikely but safe)
    extremes_df = extremes_df[~extremes_df.index.duplicated(keep='last')]
    
    return extremes_df

    # Inside plot_zigzag:
    tops_plot = [np.nan] * len(ohlc_data)
    bottoms_plot = [np.nan] * len(ohlc_data)
    marker_offset_factor = 0.01 # Small offset for visibility
    
    for idx, row in extremes.iterrows(): # Loop through confirmed extremes
         if idx < len(ohlc_data): # Check index bounds
             if row['type'] == 1: # Top
                 # Place marker slightly above the high
                 tops_plot[idx] = ohlc_data['High'].iloc[idx] * (1 + marker_offset_factor)
             elif row['type'] == -1: # Bottom
                 # Place marker slightly below the low
                 bottoms_plot[idx] = ohlc_data['Low'].iloc[idx] * (1 - marker_offset_factor)

    # Inside plot_zigzag, after preparing tops_plot and bottoms_plot:
    apds = [ # List of addplot dictionaries
        mpf.make_addplot(tops_plot, type='scatter', marker='v', markersize=50, color='green', panel=0), # Tops
        mpf.make_addplot(bottoms_plot, type='scatter', marker='^', markersize=50, color='red', panel=0)   # Bottoms
    ]

    # End of plot_zigzag function:
    mpf.plot(ohlc_data, # Main OHLC data
             type='candle', # Candlestick chart
             style='yahoo', # Visual style
             title=title,
             ylabel='Price',
             volume='Volume' in ohlc_data.columns, # Show volume panel if data exists
             addplot=apds, # Overlay the scatter markers for extremes
             figsize=(10, 6) # Suggested figure size
            )

if __name__ == '__main__':
    # --- Configuration ---
    TICKER = 'BTC-USD'
    PERIOD = '1mo'
    SIGMA = 0.03

    # --- Download Data ---
    # (Code shown previously)
    data = yf.download(...)

    # --- Calculate Extremes ---
    logging.info(f"Calculating Zigzag extremes with sigma={SIGMA:.2%}")
    extremes_df = find_zigzag_extremes(
        data['High'].to_numpy(), # Pass price arrays
        data['Low'].to_numpy(),
        data['Close'].to_numpy(),
        SIGMA
    )

    # --- Plot Results ---
    if not extremes_df.empty:
        logging.info(f"Found {len(extremes_df)} extremes.")
        logging.info("Generating plot...")
        plot_title = f"{TICKER} Price with Zigzag Extremes (Sigma={SIGMA:.1%})"
        plot_zigzag(data, extremes_df, plot_title) # Call plotting function
        logging.info("Plot generation complete.")
    else:
        logging.warning("No extremes found to plot.")
        # Optionally plot just price data
        mpf.plot(data, type='candle', title=f"{TICKER} Price Data (No Extremes Found)", ...)