"""

Convert wind time-series CSV data into a PyWake-style
directional Weibull wind resource and calculate wind-farm AEP.

Required CSV columns:
    timestamp
    wind_speed_ms
    wind_direction_deg

Example:
    timestamp,wind_speed_ms,wind_direction_deg
    2025-01-01 00:00,5.2,231
    2025-01-01 01:00,6.1,245
    2025-01-01 02:00,7.3,252

Required packages:
    numpy
    pandas
    matplotlib
    py_wake
"""


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


# ============================================================
# USER CONFIGURATION
# ============================================================

CSV_FILE = "wind_data.csv"

# Number of directional sectors used for Weibull fitting
N_DIRECTION_SECTORS = 16

# Wind-speed threshold used to identify calm observations.
# Weibull fitting is performed only on speeds above this value.
CALM_THRESHOLD = 0.20      # m/s

# Turbulence intensity used by PyWake.
# Replace this with measured TI if available.
TURBULENCE_INTENSITY = 0.10

# PyWake wind farm model
# NOJ is used here as a simple demonstration wake model.

# Sample 10-turbine layout
N_TURBINES = 10

# Turbine spacing
TURBINE_SPACING = 500.0    # m

# Output files
RESOURCE_OUTPUT_CSV = "directional_weibull_parameters.csv"

# Plot settings
MAKE_PLOTS = True


# ============================================================
# READ WIND TIME-SERIES DATA
# ============================================================

def read_wind_data(filename):
    """
    Read wind-speed and wind-direction time-series data.

    Expected columns:
        timestamp
        wind_speed_ms
        wind_direction_deg
    """

    if not os.path.isfile(filename):
        print(
            f"CSV file not found: {filename}\n"
        )
        exit(0)
        
    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"Missing required column: {column}"
            )

    # Convert to numeric
    df["wind_speed_ms"] = pd.to_numeric(
        df["wind_speed_ms"],
        errors="coerce"
    )

    df["wind_direction_deg"] = pd.to_numeric(
        df["wind_direction_deg"],
        errors="coerce"
    )

    # Remove invalid observations
    valid = (
        np.isfinite(df["wind_speed_ms"])
        &
        np.isfinite(df["wind_direction_deg"])
        &
        (df["wind_speed_ms"] >= 0)
    )

    df = df.loc[valid].copy()

    # Normalize direction to 0-360 degrees
    df["wind_direction_deg"] = (
        df["wind_direction_deg"] % 360
    )

    if len(df) == 0:
        raise ValueError(
            "No valid wind observations found."
        )

    return df


# ============================================================
# CREATE DIRECTIONAL SECTORS
# ============================================================

def assign_direction_sector(
        direction,
        n_sectors):
    """
    Assign each wind direction to a sector.

    Example for 16 sectors:

        0
        22.5
        45
        ...

    Sector centers are:

        0, 22.5, 45, ..., 337.5 degrees

    The sectors are centered on these directions.
    """

    sector_width = 360.0 / n_sectors

    # Shift by half a sector so that
    # 0 degrees lies at the center of a sector.
    sector_index = (
        np.floor(
            (direction + sector_width / 2)
            / sector_width
        ).astype(int)
        % n_sectors
    )

    return sector_index


# ============================================================
# ESTIMATE WEIBULL MLE
# ============================================================

def estimate_weibull_mle(wind_speed):
    """
    Estimate Weibull shape k and scale A using MLE

    Weibull PDF:

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

    The MLE equation for k is solved using
    Newton-Raphson iteration.

    Returns
    -------
    k : float
        Weibull shape parameter.

    A : float
        Weibull scale parameter [m/s].
    """

    v = np.asarray(wind_speed, dtype=float)

    # Weibull MLE requires positive values
    v = v[v > CALM_THRESHOLD]

    if len(v) < 10:
        return np.nan, np.nan

    log_v = np.log(v)

    # Initial estimate
    k = 2.0

    for _ 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)

        # MLE equation
        f = (
            1.0 / k
            + mean_log_v
            - sum_vk_logv / sum_vk
        )

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

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

        if abs(derivative) < 1e-12:
            break

        k_new = k - f / derivative

        # Prevent numerical problems
        if k_new <= 0:
            k_new = k / 2

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

        k = k_new

    # Weibull scale parameter
    A = (
        np.mean(v ** k)
    ) ** (1.0 / k)

    return k, A


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

def weibull_pdf(v, k, A):
    """
    Calculate Weibull probability density
    """

    v = np.asarray(v)

    pdf = np.zeros_like(
        v,
        dtype=float
    )

    positive = v > 0

    pdf[positive] = (
        (k / A)
        *
        (v[positive] / A) ** (k - 1)
        *
        np.exp(
            -(v[positive] / A) ** k
        )
    )

    return pdf


# ============================================================
# CREATE DIRECTIONAL WEIBULL RESOURCE
# ============================================================

def create_directional_weibull_resource(
        df,
        n_sectors):
    """
    Calculate:

        directional frequency
        Weibull k
        Weibull A

    for every direction sector.
    """

    df = df.copy()

    # Remove calm observations from Weibull resource
    df = df[
        df["wind_speed_ms"]
        > CALM_THRESHOLD
    ].copy()

    if len(df) == 0:
        raise ValueError(
            "No wind-speed observations above "
            "the calm threshold."
        )

    # Assign sectors
    df["sector"] = assign_direction_sector(
        df["wind_direction_deg"].values,
        n_sectors
    )

    sector_width = (
        360.0 / n_sectors
    )

    sector_centers = (
        np.arange(n_sectors)
        * sector_width
    )

    # Output arrays
    frequency = np.zeros(n_sectors)
    weibull_k = np.full(
        n_sectors,
        np.nan
    )
    weibull_A = np.full(
        n_sectors,
        np.nan
    )

    # --------------------------------------------------------
    # Frequency
    # --------------------------------------------------------

    sector_counts = (
        df["sector"]
        .value_counts()
        .reindex(
            range(n_sectors),
            fill_value=0
        )
    )

    frequency = (
        sector_counts.values
        /
        sector_counts.values.sum()
    )

    # --------------------------------------------------------
    # Weibull parameters
    # --------------------------------------------------------

    for sector in range(n_sectors):

        sector_speed = df.loc[
            df["sector"] == sector,
            "wind_speed_ms"
        ].values

        if len(sector_speed) >= 10:

            k, A = estimate_weibull_mle(
                sector_speed
            )

            weibull_k[sector] = k
            weibull_A[sector] = A

    # --------------------------------------------------------
    # Deal with sectors with insufficient data
    # --------------------------------------------------------

    valid = (
        np.isfinite(weibull_k)
        &
        np.isfinite(weibull_A)
        &
        (frequency > 0)
    )

    if not np.all(valid):

        print(
            "\nWARNING:"
            "\nSome direction sectors do not have "
            "enough data for a Weibull fit."
        )

        print(
            "Affected sectors:",
            sector_centers[~valid]
        )

    # For a production PyWake model, every sector with
    # non-zero probability needs valid Weibull parameters.
    if np.any(
        (~np.isfinite(weibull_k))
        &
        (frequency > 0)
    ):
        raise ValueError(
            "One or more populated direction sectors "
            "do not have enough observations for "
            "Weibull fitting."
        )

    # Normalize frequency
    frequency /= frequency.sum()

    return (
        sector_centers,
        frequency,
        weibull_k,
        weibull_A
    )


# ============================================================
# SAVE WIND RESOURCE
# ============================================================

def save_resource(
        filename,
        directions,
        frequency,
        weibull_k,
        weibull_A):
    """
    Save PyWake-style directional resource parameters.
    """

    output = pd.DataFrame({

        "wind_direction_deg":
            directions,

        "sector_frequency":
            frequency,

        "weibull_k":
            weibull_k,

        "weibull_A_ms":
            weibull_A
    })

    output.to_csv(
        filename,
        index=False
    )

    print(
        f"\nResource parameters written to:"
        f" {filename}"
    )


# ============================================================
# CREATE PYWAKE SITE
# ============================================================

def create_pywake_site(
        frequency,
        weibull_A,
        weibull_k,
        turbulence_intensity):
    """
    Create a PyWake UniformWeibullSite.

    PyWake expects:

        p_wd = sector frequency
        a    = Weibull scale
        k    = Weibull shape
    """

    try:

        from py_wake.site import (
            UniformWeibullSite
        )

    except ImportError:

        raise ImportError(
            "\nPyWake is not installed.\n"
            "Install PyWake before running "
            "the wake calculation."
        )

    site = UniformWeibullSite(

        p_wd=frequency,

        a=weibull_A,

        k=weibull_k,

        ti=turbulence_intensity,

        # Nearest is generally appropriate for
        # discrete sector data.
        interp_method="nearest"
    )

    return site


# ============================================================
# CREATE SAMPLE 10-TURBINE LAYOUT
# ============================================================

def create_sample_layout():
    """
    Create a simple 10-turbine staggered layout.

    The layout is only an example.

    Replace this function with your optimizer-generated
    turbine coordinates for the actual wind-farm problem.
    """

    spacing = TURBINE_SPACING

    x = np.array([
        0,
        1 * spacing,
        2 * spacing,
        3 * spacing,
        4 * spacing,

        0.5 * spacing,
        1.5 * spacing,
        2.5 * spacing,
        3.5 * spacing,
        4.5 * spacing
    ])

    y = np.array([
        0,
        0,
        0,
        0,
        0,

        0.8 * spacing,
        0.8 * spacing,
        0.8 * spacing,
        0.8 * spacing,
        0.8 * spacing
    ])

    return x, y


# ============================================================
# RUN PYWAKE
# ============================================================

def run_pywake(
        site,
        x,
        y):
    """
    Run a PyWake NOJ wake calculation using
    the Vestas V80 example turbine.

    Returns
    -------
    simulation_result
    """

    try:

        from py_wake import NOJ

        from py_wake.examples.data.hornsrev1 import (
            V80
        )

    except ImportError:

        raise ImportError(
            "PyWake installation/import failed."
        )

    # V80 is a PyWake example turbine.
    wind_turbines = V80()

    # Create NOJ wake model
    wind_farm_model = NOJ(
        site,
        wind_turbines
    )

    # Run simulation
    simulation_result = wind_farm_model(
        x,
        y
    )

    return simulation_result


# ============================================================
# CALCULATE AEP
# ============================================================

def calculate_aep(
        simulation_result):
    """
    Calculate:

        AEP with wake
        AEP without wake
        wake loss
    """

    # With wake
    aep_with_wake = (
        simulation_result
        .aep(with_wake_loss=True)
        .sum()
        .item()
    )

    # Without wake
    aep_without_wake = (
        simulation_result
        .aep(with_wake_loss=False)
        .sum()
        .item()
    )

    wake_loss = (
        (
            aep_without_wake
            -
            aep_with_wake
        )
        /
        aep_without_wake
        *
        100
    )

    return (
        aep_with_wake,
        aep_without_wake,
        wake_loss
    )


# ============================================================
# PLOT WIND ROSE
# ============================================================

def plot_wind_rose(
        df,
        n_sectors):
    """
    Plot a simple directional wind rose.
    """

    speed = df["wind_speed_ms"].values
    direction = df["wind_direction_deg"].values

    # Remove calm
    valid = speed > CALM_THRESHOLD

    speed = speed[valid]
    direction = direction[valid]

    sector_width = (
        360.0 / n_sectors
    )

    sector_index = assign_direction_sector(
        direction,
        n_sectors
    )

    # Speed bins
    speed_bins = np.arange(
        0,
        np.ceil(np.max(speed)) + 2,
        2
    )

    counts = np.zeros(
        (
            n_sectors,
            len(speed_bins) - 1
        )
    )

    speed_index = (
        np.digitize(
            speed,
            speed_bins
        ) - 1
    )

    for d, s in zip(
        sector_index,
        speed_index
    ):

        if (
            0 <= d < n_sectors
            and
            0 <= s < len(speed_bins) - 1
        ):
            counts[d, s] += 1

    counts /= counts.sum()
    counts *= 100

    direction_centers = (
        np.arange(n_sectors)
        * sector_width
    )

    theta = np.deg2rad(
        direction_centers
    )

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

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

    bottom = np.zeros(n_sectors)

    for i in range(
        len(speed_bins) - 1
    ):

        values = counts[:, i]

        ax.bar(
            theta,
            values,
            width=np.deg2rad(
                sector_width
            ),
            bottom=bottom,
            align="center",
            edgecolor="black",
            linewidth=0.4,
            alpha=0.8,
            label=(
                f"{speed_bins[i]:.0f}-"
                f"{speed_bins[i+1]:.0f} m/s"
            )
        )

        bottom += values

    ax.set_theta_zero_location("N")
    ax.set_theta_direction(-1)

    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=20
    )

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

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


# ============================================================
# PLOT WEIBULL PARAMETERS
# ============================================================

def plot_weibull_parameters(
        directions,
        frequency,
        weibull_k,
        weibull_A):
    """
    Plot directional Weibull k and A.
    """

    fig, axes = plt.subplots(
        2,
        1,
        figsize=(10, 8),
        sharex=True
    )

    axes[0].bar(
        directions,
        frequency * 100,
        width=360 / len(directions) * 0.8
    )

    axes[0].set_ylabel(
        "Frequency (%)"
    )

    axes[0].set_title(
        "Directional Wind Resource"
    )

    axes[0].grid(
        True,
        alpha=0.3
    )

    axes[1].plot(
        directions,
        weibull_A,
        marker="o",
        label="Weibull A"
    )

    axes[1].plot(
        directions,
        weibull_k,
        marker="s",
        label="Weibull k"
    )

    axes[1].set_xlabel(
        "Wind direction (deg)"
    )

    axes[1].set_ylabel(
        "Parameter value"
    )

    axes[1].grid(
        True,
        alpha=0.3
    )

    axes[1].legend()

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


# ============================================================
# PLOT WEIBULL FITS
# ============================================================

def plot_directional_weibull_fits(
        df,
        directions,
        weibull_k,
        weibull_A):
    """
    Plot measured and fitted Weibull distributions
    for each direction sector.
    """

    n_sectors = len(directions)

    n_columns = 4

    n_rows = int(
        np.ceil(
            n_sectors / n_columns
        )
    )

    fig, axes = plt.subplots(
        n_rows,
        n_columns,
        figsize=(14, 3 * n_rows)
    )

    axes = np.asarray(
        axes
    ).reshape(-1)

    df = df.copy()

    df["sector"] = assign_direction_sector(
        df["wind_direction_deg"].values,
        n_sectors
    )

    for sector in range(n_sectors):

        ax = axes[sector]

        speeds = df.loc[
            df["sector"] == sector,
            "wind_speed_ms"
        ].values

        speeds = speeds[
            speeds > CALM_THRESHOLD
        ]

        if len(speeds) < 10:
            ax.set_visible(False)
            continue

        ax.hist(
            speeds,
            bins=15,
            density=True,
            alpha=0.5
        )

        x = np.linspace(
            0.01,
            max(
                np.percentile(
                    speeds,
                    99
                ),
                1
            ),
            300
        )

        pdf = weibull_pdf(
            x,
            weibull_k[sector],
            weibull_A[sector]
        )

        ax.plot(
            x,
            pdf,
            linewidth=2
        )

        ax.set_title(
            f"{directions[sector]:.1f}°\n"
            f"k={weibull_k[sector]:.2f}, "
            f"A={weibull_A[sector]:.2f}"
        )

        ax.set_xlabel(
            "Wind speed (m/s)"
        )

        ax.set_ylabel(
            "PDF"
        )

        ax.grid(
            True,
            alpha=0.3
        )

    # Hide unused axes
    for i in range(
        n_sectors,
        len(axes)
    ):
        axes[i].set_visible(False)

    fig.suptitle(
        "Directional Weibull Fits",
        fontsize=16
    )

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


# ============================================================
# PLOT TURBINE LAYOUT
# ============================================================

def plot_layout(
        x,
        y):
    """
    Plot sample wind farm layout.
    """

    plt.figure(
        figsize=(10, 5)
    )

    plt.scatter(
        x,
        y,
        s=150
    )

    for i in range(len(x)):

        plt.text(
            x[i] + 15,
            y[i] + 15,
            f"T{i + 1}"
        )

    plt.xlabel(
        "x (m)"
    )

    plt.ylabel(
        "y (m)"
    )

    plt.title(
        "10-Turbine Sample Layout"
    )

    plt.axis("equal")
    plt.grid(True)

    plt.tight_layout()
    plt.savefig("Wind-Turbine-Layout.png")
    #plt.show()


# ============================================================
# DEFINE MAIN FUNCTION: ENTRY POINT
# ============================================================

def main():

    print("=" * 65)
    print("WIND RESOURCE -> PYWAKE -> AEP")
    print("=" * 65)

    # --------------------------------------------------------
    # Read time-series data
    # --------------------------------------------------------

    print("\nReading wind data...")

    df = read_wind_data(
        CSV_FILE
    )

    print(
        f"Valid observations: {len(df)}"
    )

    # --------------------------------------------------------
    # Basic statistics
    # --------------------------------------------------------

    wind_speed = (
        df["wind_speed_ms"]
        .values
    )

    calm_count = np.sum(
        wind_speed <= CALM_THRESHOLD
    )

    calm_fraction = (
        calm_count
        /
        len(wind_speed)
        *
        100
    )

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

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

    print(
        f"Calm/near-zero fraction: "
        f"{calm_fraction:.2f}%"
    )

    # --------------------------------------------------------
    # Directional Weibull analysis
    # --------------------------------------------------------

    print(
        "\nCalculating directional Weibull parameters..."
    )

    (
        directions,
        frequency,
        weibull_k,
        weibull_A
    ) = create_directional_weibull_resource(
        df,
        N_DIRECTION_SECTORS
    )

    # --------------------------------------------------------
    # Print resource table
    # --------------------------------------------------------

    resource_table = pd.DataFrame({

        "Direction_deg":
            directions,

        "Frequency":
            frequency,

        "Frequency_%":
            frequency * 100,

        "Weibull_k":
            weibull_k,

        "Weibull_A_ms":
            weibull_A
    })

    print("\nDirectional wind resource")
    print("-" * 65)

    print(
        resource_table.to_string(
            index=False,
            float_format=lambda x:
                f"{x:.4f}"
        )
    )

    print(
        "\nFrequency sum:",
        frequency.sum()
    )

    # --------------------------------------------------------
    # Save resource
    # --------------------------------------------------------

    save_resource(
        RESOURCE_OUTPUT_CSV,
        directions,
        frequency,
        weibull_k,
        weibull_A
    )

    # --------------------------------------------------------
    # Create PyWake site
    # --------------------------------------------------------

    print(
        "\nCreating PyWake UniformWeibullSite..."
    )

    site = create_pywake_site(
        frequency,
        weibull_A,
        weibull_k,
        TURBULENCE_INTENSITY
    )

    print(
        "PyWake site created successfully."
    )

    # --------------------------------------------------------
    # Create sample layout
    # --------------------------------------------------------

    x, y = create_sample_layout()

    print(
        f"\nNumber of turbines: {len(x)}"
    )

    # --------------------------------------------------------
    # Run PyWake
    # --------------------------------------------------------

    print(
        "\nRunning PyWake NOJ wake calculation..."
    )

    simulation_result = run_pywake(
        site,
        x,
        y
    )

    print(
        "PyWake simulation completed."
    )

    # --------------------------------------------------------
    # AEP
    # --------------------------------------------------------

    (
        aep_with_wake,
        aep_without_wake,
        wake_loss
    ) = calculate_aep(
        simulation_result
    )

    print("\n" + "=" * 65)
    print("WIND FARM RESULTS")
    print("=" * 65)

    print(
        f"AEP without wake : "
        f"{aep_without_wake:.3f} GWh/year"
    )

    print(
        f"AEP with wake    : "
        f"{aep_with_wake:.3f} GWh/year"
    )

    print(
        f"Wake loss        : "
        f"{wake_loss:.2f} %"
    )

    print(
        f"Average AEP/turbine: "
        f"{aep_with_wake / N_TURBINES:.3f} "
        f"GWh/year"
    )

    # --------------------------------------------------------
    # Capacity factor
    #
    # V80 nominal rating is approximately 2 MW.
    # Use the actual turbine rating for your turbine model.
    # --------------------------------------------------------

    RATED_POWER_MW = 2.0

    theoretical_max_gwh = (
        N_TURBINES
        *
        RATED_POWER_MW
        *
        8760
        /
        1000
    )

    capacity_factor = (
        aep_with_wake
        /
        theoretical_max_gwh
        *
        100
    )

    print(
        f"Capacity factor  : "
        f"{capacity_factor:.2f} %"
    )

    print("=" * 65)

    # --------------------------------------------------------
    # Plots
    # --------------------------------------------------------

    if MAKE_PLOTS:

        plot_wind_rose(
            df,
            N_DIRECTION_SECTORS
        )

        plot_weibull_parameters(
            directions,
            frequency,
            weibull_k,
            weibull_A
        )

        plot_directional_weibull_fits(
            df,
            directions,
            weibull_k,
            weibull_A
        )

        plot_layout(
            x,
            y
        )


# ============================================================
# PROGRAM ENTRY POINT
# ============================================================

if __name__ == "__main__":
    main()