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

"""
  Creates a sample wind time-series CSV

  Reads wind speed and direction from the CSV

  Creates a wind rose from first principles using only NumPy + 
  Matplotlib

  Estimates Weibull shape k and scale A using the Maximum Likelihood 
  Estimation (MLE) approach

  Plots the fitted Weibull distribution against the measured wind-
  speed histogram

  Prints useful wind statistics.
  
Note: 
This code creates a single Weibull distribution for the complete time
series ignores directional dependence. For a wind-farm yield model, it 
is often more useful to estimate a Weibull distribution for each wind-
direction sector, then combine those sector-wise Weibull distributions 
with the directional frequencies.
"""
# =====================================================================
# USER SETTINGS
# =====================================================================

CSV_FILE = "sample_wind_data.csv"

# Wind rose settings
N_DIRECTION_SECTORS = 16
N_SPEED_BINS = 8

# Weibull fitting
MIN_WIND_SPEED = 0.5     # Ignore calm/near-zero observations

# =====================================================================
# READ WIND TIME SERIES
# =====================================================================

def read_wind_data(filename):
    """
    Read wind data from CSV
    """

    df = pd.read_csv(filename)

    required_columns = [
        "wind_speed_ms",
        "wind_direction_deg"
    ]

    for column in required_columns:
        if column not in df.columns:
            raise ValueError(
                f"CSV must contain column: {column}"
            )

    speed = df["wind_speed_ms"].to_numpy(dtype=float)
    direction = df["wind_direction_deg"].to_numpy(dtype=float)

    # Remove invalid observations
    valid = (
        np.isfinite(speed)
        & np.isfinite(direction)
        & (speed >= 0)
    )

    speed = speed[valid]
    direction = direction[valid] % 360

    return speed, direction


# =====================================================================
# CREATE WIND ROSE FROM FIRST PRINCIPLES
# =====================================================================

def calculate_wind_rose(
        wind_speed,
        wind_direction,
        n_direction_sectors=16,
        speed_bins=None):
    """
    Calculate frequency of wind speeds for each direction sector.

    Returns:
        frequency_matrix
        direction_centers
        speed_bin_edges
    """

    if speed_bins is None:
        max_speed = np.ceil(np.max(wind_speed))
        speed_bins = np.linspace(0, max_speed, N_SPEED_BINS + 1)

    direction_sector_width = 360 / n_direction_sectors

    # Convert direction to sector number
    sector_index = (
        np.floor(
            wind_direction / direction_sector_width
        ).astype(int)
        % n_direction_sectors
    )

    # Speed bin number
    speed_index = (
        np.digitize(wind_speed, speed_bins) - 1
    )

    n_speed_bins = len(speed_bins) - 1

    frequency_matrix = np.zeros(
        (n_direction_sectors, n_speed_bins)
    )

    # Count observations
    for d, s in zip(sector_index, speed_index):

        if 0 <= s < n_speed_bins:
            frequency_matrix[d, s] += 1

    # Convert counts to percentage of all observations
    frequency_matrix = (
        frequency_matrix
        / len(wind_speed)
        * 100
    )

    direction_centers = (
        np.arange(n_direction_sectors)
        * direction_sector_width
        + direction_sector_width / 2
    )

    return (
        frequency_matrix,
        direction_centers,
        speed_bins
    )


def plot_wind_rose(
        frequency_matrix,
        direction_centers,
        speed_bins):
    """
    Plot stacked wind rose using Matplotlib polar axes.
    """

    n_direction_sectors = len(direction_centers)

    sector_width = (
        2 * np.pi / n_direction_sectors
    )

    # Matplotlib polar coordinates:
    # 0 rad = East
    # pi/2  = North
    #
    # Convert meteorological direction:
    # 0 deg  = North
    # 90 deg = East
    #
    theta = np.deg2rad(
        90 - direction_centers
    )

    fig = plt.figure(figsize=(9, 9))

    ax = fig.add_subplot(
        111,
        polar=True
    )

    bottom = np.zeros(n_direction_sectors)

    for speed_bin in range(
            frequency_matrix.shape[1]):

        values = frequency_matrix[:, speed_bin]

        ax.bar(
            theta,
            values,
            width=sector_width,
            bottom=bottom,
            align="center",
            edgecolor="black",
            linewidth=0.5,
            alpha=0.8,
            label=(
                f"{speed_bins[speed_bin]:.1f}"
                f"–"
                f"{speed_bins[speed_bin + 1]:.1f} m/s"
            )
        )

        bottom += values

    # Put North at the top
    ax.set_theta_zero_location("N")

    # Clockwise direction
    ax.set_theta_direction(-1)

    # Direction labels
    ax.set_xticks(
        np.deg2rad(
            np.arange(0, 360, 45)
        )
    )

    ax.set_xticklabels([
        "N", "NE", "E", "SE",
        "S", "SW", "W", "NW"
    ])

    ax.set_title(
        "Wind Rose",
        pad=25,
        fontsize=16
    )

    ax.set_ylabel(
        "Frequency (%)",
        labelpad=30
    )

    ax.legend(
        title="Wind speed",
        loc="upper left",
        bbox_to_anchor=(1.10, 1.05)
    )

    plt.tight_layout()
    plt.savefig("Wind-Rose-Diagram.png")
    #plt.show()


# =====================================================================
# WEIBULL MLE ESTIMATION
# =====================================================================

def estimate_weibull_mle(wind_speed):
    """
    Estimate Weibull shape parameter k and scale parameter A
    using Maximum Likelihood Estimation.

    Weibull PDF:

        f(v) = (k/A)
               (v/A)^(k-1)
               exp[-(v/A)^k]

    For a given k:

        A = [mean(v^k)]^(1/k)

    k is solved iteratively from:

        1/k + mean(ln(v))
        - sum(v^k ln(v)) / sum(v^k) = 0
    """

    # Remove zero and negative values
    v = wind_speed[
        wind_speed > MIN_WIND_SPEED
    ]

    if len(v) < 10:
        raise ValueError(
            "Not enough positive wind-speed observations."
        )

    log_v = np.log(v)

    # Initial guess
    k = 2.0

    # Newton-Raphson iteration
    for iteration in range(100):

        v_k = v ** k

        sum_vk = np.sum(v_k)

        sum_vk_logv = np.sum(
            v_k * log_v
        )

        mean_log_v = np.mean(log_v)

        # Function
        f = (
            1 / k
            + mean_log_v
            - sum_vk_logv / sum_vk
        )

        # Derivative
        sum_vk_logv2 = np.sum(
            v_k * log_v ** 2
        )

        derivative = (
            -1 / k**2
            -
            (
                sum_vk_logv2 * sum_vk
                - sum_vk_logv**2
            )
            / sum_vk**2
        )

        k_new = k - f / derivative

        if abs(k_new - k) < 1e-8:
            k = k_new
            break

        k = k_new

    # Scale parameter
    A = (
        np.mean(v ** k)
    ) ** (1 / k)

    return k, A


# =====================================================================
# WEIBULL PDF
# =====================================================================

def weibull_pdf(wind_speed, k, A):
    """
    Weibull probability density function.
    """

    return (
        (k / A)
        * (wind_speed / A) ** (k - 1)
        * np.exp(
            -(wind_speed / A) ** k
        )
    )


# =====================================================================
# PLOT WEIBULL DISTRIBUTION
# =====================================================================

def plot_weibull_distribution(
        wind_speed,
        k,
        A):
    """
    Plot measured wind-speed distribution
    and fitted Weibull distribution.
    """

    max_speed = np.percentile(
        wind_speed,
        99.5
    )

    x = np.linspace(
        0.01,
        max_speed,
        500
    )

    pdf = weibull_pdf(
        x,
        k,
        A
    )

    plt.figure(figsize=(9, 6))

    # Measured distribution
    plt.hist(
        wind_speed,
        bins=30,
        density=True,
        alpha=0.5,
        label="Measured wind speed"
    )

    # Weibull PDF
    plt.plot(
        x,
        pdf,
        linewidth=2,
        label=(
            f"Weibull fit: "
            f"k = {k:.3f}, A = {A:.3f} m/s"
        )
    )

    plt.xlabel(
        "Wind speed (m/s)"
    )

    plt.ylabel(
        "Probability density"
    )

    plt.title(
        "Wind Speed Distribution and Weibull Fit"
    )

    plt.grid(
        True,
        alpha=0.3
    )

    plt.legend()

    plt.tight_layout()
    plt.savefig("Wind-Weibull-Curve.png")
    #plt.show()


# =====================================================================
# MAIN PROGRAM
# =====================================================================

def main():

    # --------------------------------------------------------
    # Read CSV
    # --------------------------------------------------------

    wind_speed, wind_direction = (
        read_wind_data(CSV_FILE)
    )

    print("\nWind data")
    print("-" * 40)

    print(
        f"Number of observations : "
        f"{len(wind_speed)}"
    )

    print(
        f"Mean wind speed        : "
        f"{np.mean(wind_speed):.3f} m/s"
    )

    print(
        f"Maximum wind speed     : "
        f"{np.max(wind_speed):.3f} m/s"
    )

    print(
        f"Median wind speed      : "
        f"{np.median(wind_speed):.3f} m/s"
    )

    # --------------------------------------------------------
    # Wind rose
    # --------------------------------------------------------

    frequency_matrix, direction_centers, speed_bins = (
        calculate_wind_rose(
            wind_speed,
            wind_direction,
            n_direction_sectors=N_DIRECTION_SECTORS
        )
    )

    plot_wind_rose(
        frequency_matrix,
        direction_centers,
        speed_bins
    )

    # --------------------------------------------------------
    # Weibull estimation
    # --------------------------------------------------------

    k, A = estimate_weibull_mle(
        wind_speed
    )

    print("\nWeibull parameters")
    print("-" * 40)

    print(
        f"Shape parameter k : {k:.4f}"
    )

    print(
        f"Scale parameter A : {A:.4f} m/s"
    )

    # --------------------------------------------------------
    # Derived quantities
    # --------------------------------------------------------

    # Mean wind speed from Weibull distribution
    mean_weibull = (
        A
        * np.math.gamma(
            1 + 1 / k
        )
    )

    # Most probable wind speed
    most_probable_speed = (
        A
        * ((k - 1) / k) ** (1 / k)
        if k > 1
        else 0
    )

    # Wind speed carrying maximum energy
    energy_speed = (
        A
        * ((k + 2) / k) ** (1 / k)
    )

    print("\nDerived Weibull statistics")
    print("-" * 40)

    print(
        f"Mean speed from Weibull : "
        f"{mean_weibull:.3f} m/s"
    )

    print(
        f"Most probable speed     : "
        f"{most_probable_speed:.3f} m/s"
    )

    print(
        f"Energy-carrying speed   : "
        f"{energy_speed:.3f} m/s \n"
    )

    # --------------------------------------------------------
    # Weibull plot
    # --------------------------------------------------------

    plot_weibull_distribution(
        wind_speed,
        k,
        A
    )


if __name__ == "__main__":
    main()