import numpy as np
import pandas as pd
import pvlib

def night_mask(
    df,
    lat=None,
    lon=None,
    time_col="time",
    elev_min=0.0,
    ssrd_col=None,
    hour_window=None,
):
    """
    判断夜间点，返回 bool 数组：
    True  = 夜间
    False = 白天

    优先级：
    1. lat/lon 均存在 -> 根据太阳高度角判断
    2. ssrd_col 存在 -> 辐照 <= 0 判为夜间
    3. hour_window 存在 -> 窗口外判为夜间
    """

    # 1. 根据太阳高度角判断
    times = pd.to_datetime(df[time_col], errors="coerce")

    # 如果时间没有时区，这里按 UTC 处理
    if times.dt.tz is None:
        times = times.dt.tz_localize("UTC")
    else:
        times = times.dt.tz_convert("UTC")

    solar_position = pvlib.solarposition.get_solarposition(
        time=times,
        latitude=float(lat),
        longitude=float(lon),
    )

    elevation = solar_position["apparent_elevation"].to_numpy()

    return elevation <= float(elev_min)

    # 2. 没有经纬度时，可根据太阳辐照判断
    if ssrd_col is not None and ssrd_col in df.columns:
        ssrd = pd.to_numeric(df[ssrd_col], errors="coerce")
        return (ssrd <= 0).fillna(False).to_numpy()

    # 3. 根据小时范围粗略判断
    if hour_window is not None:
        start_hour, end_hour = hour_window

        times = pd.to_datetime(df[time_col], errors="coerce")
        hours = times.dt.hour

        if start_hour <= end_hour:
            daytime = (hours >= start_hour) & (hours <= end_hour)
        else:
            daytime = (hours >= start_hour) | (hours <= end_hour)

        return (~daytime).fillna(False).to_numpy()

    raise ValueError(
        "无法判断昼夜：请提供 lat/lon、ssrd_col 或 hour_window"
    )


def solor_processed_power(df, n, power_col="v1-Dlinear-288", *, lat=None, lon=None,
                      elev_min=0.0, clip_negative=True, ssrd_col=None,
                      hour_window=None, time_col="time"):
    """只对前 n 行: 夜间置 0(按太阳高度角判定) + 负值置 0。"""
    processed = df.copy()
    m = min(int(n), len(processed))
    if m <= 0:
        return processed
    head = processed.iloc[:m]
    night = night_mask(head, lat=lat, lon=lon, time_col=time_col, elev_min=elev_min,
                       ssrd_col=ssrd_col, hour_window=hour_window)
    p = pd.to_numeric(head[power_col], errors="coerce")
    bad = np.asarray(night, dtype=bool)
    if clip_negative:
        bad |= (p < 0).fillna(False).to_numpy()
    if bad.any():
        processed.loc[processed.index[:m], power_col] = p.mask(bad, 0.0).to_numpy()
    return processed