# -*-coding:utf-8 -*-
'''
@Author:  Haoran Yu
@Contact: yuhaoran251@mails.ucas.ac.cn
@Date: 2025/12/5 18:25
@Version: 2.1
@Copyright: Copyright (c) 2025 Haoran Yu
@Author: Haoran Yu
@Desc: 极端气候指标计算函数库
'''
import numpy as np
import xarray as xr
from tqdm import tqdm
import warnings
from datetime import datetime, timedelta

"""======================================================辅助函数======================================================"""
"""======================================================辅助函数======================================================"""
"""======================================================辅助函数======================================================"""
"""======================================================辅助函数======================================================"""
"""======================================================辅助函数======================================================"""
"""======================================"""
def _convert_time_to_datetime(pr):
    """
    稳健的时间转换函数：将任意格式的time维度转换为datetime64[D]（日分辨率）
    支持各种cftime日历类型、数值格式、字符串格式和datetime64格式
    自动修正无效日期（如2月30日）
    """
    time_vals = pr.time.values

    # ==================== 1. 处理cftime类型 ====================
    try:
        import cftime
        if len(time_vals) > 0 and isinstance(time_vals[0], cftime.datetime):
            parsed_times = []
            for t in time_vals:
                year, month, day = t.year, t.month, t.day

                # 修正无效日期
                if month == 2:
                    calendar = getattr(t, 'calendar', 'standard')
                    if calendar in ['noleap', '365_day', '365']:
                        day = min(day, 28)  # 无闰年日历
                    elif calendar in ['360_day', '360']:
                        day = min(day, 30)  # 360天日历
                    else:
                        # 标准日历：检查闰年
                        is_leap = (year % 4 == 0 and year % 100 != 0) or (year % 400 == 0)
                        day = min(day, 29 if is_leap else 28)
                elif month in [4, 6, 9, 11]:
                    day = min(day, 30)  # 小月
                else:
                    day = min(day, 31)  # 大月

                try:
                    dt = datetime(year, month, day)
                    parsed_times.append(np.datetime64(dt))
                except ValueError:
                    # 兜底修正
                    if month == 2:
                        day = 28
                    elif month in [4, 6, 9, 11]:
                        day = 30
                    else:
                        day = 31
                    dt = datetime(year, month, day)
                    parsed_times.append(np.datetime64(dt))

            return pr.assign_coords(time=np.array(parsed_times, dtype='datetime64[D]'))
    except ImportError:
        pass
    except Exception:
        pass  # 继续尝试其他格式

    # ==================== 2. 处理datetime64类型 ====================
    if len(time_vals) > 0 and np.issubdtype(time_vals.dtype, np.datetime64):
        return pr.assign_coords(time=pr.time.dt.floor('D'))

    # ==================== 3. 处理数值型时间 ====================
    if len(time_vals) > 0 and np.issubdtype(time_vals.dtype, np.number):
        parsed_times = []
        for val in time_vals:
            val = float(val)

            # YYYYMMDD格式
            if 19000101 <= val <= 21001231:
                val_int = int(val)
                year = val_int // 10000
                month = (val_int // 100) % 100
                day = val_int % 100

                # 修正日期
                if month == 2:
                    is_leap = (year % 4 == 0 and year % 100 != 0) or (year % 400 == 0)
                    day = min(day, 29 if is_leap else 28)
                elif month in [4, 6, 9, 11]:
                    day = min(day, 30)
                else:
                    day = min(day, 31)

                dt = datetime(year, month, day)
                parsed_times.append(np.datetime64(dt))

            # 年+日序格式
            elif 1900 <= val < 2100:
                year = int(val)
                doy = int(round((val - year) * 1000))
                doy = max(1, min(doy, 366))
                dt = datetime(year, 1, 1) + timedelta(days=doy - 1)
                parsed_times.append(np.datetime64(dt))

            # 时间戳
            elif val > 1e9:
                if val > 1e12:
                    val = val / 1000
                dt = datetime.fromtimestamp(val)
                parsed_times.append(np.datetime64(dt))

        if parsed_times:
            return pr.assign_coords(time=np.array(parsed_times, dtype='datetime64[D]'))

    # ==================== 4. 处理字符串型时间 ====================
    if len(time_vals) > 0 and not np.issubdtype(time_vals.dtype, np.number):
        parsed_times = []
        date_formats = [
            '%Y-%m-%d %H:%M:%S', '%Y-%m-%d %H:%M:%S.%f',
            '%Y-%m-%d %H:%M', '%Y-%m-%d',
            '%Y/%m/%d', '%Y%m%d', '%Y.%m.%d'
        ]

        for date_str in [str(t).strip() for t in time_vals]:
            parsed = None

            # 尝试各种格式
            for fmt in date_formats:
                try:
                    parsed = datetime.strptime(date_str, fmt)
                    break
                except ValueError as e:
                    # 尝试修正无效日期
                    if "day is out of range for month" in str(e):
                        try:
                            date_part = date_str.split(' ')[0] if ' ' in date_str else date_str
                            parts = date_part.replace('/', '-').replace('.', '-').split('-')

                            if len(parts) == 3:
                                if len(parts[0]) == 4:  # YYYY-MM-DD
                                    year, month, day = map(int, parts)
                                else:  # 尝试其他顺序
                                    try:
                                        datetime.strptime(date_part, '%d-%m-%Y')
                                        day, month, year = map(int, parts)
                                    except:
                                        continue

                                # 修正日期
                                if month == 2:
                                    is_leap = (year % 4 == 0 and year % 100 != 0) or (year % 400 == 0)
                                    day = min(day, 29 if is_leap else 28)
                                elif month in [4, 6, 9, 11]:
                                    day = min(day, 30)
                                else:
                                    day = min(day, 31)

                                corrected_date = f"{year:04d}-{month:02d}-{day:02d}"
                                if ' ' in date_str:
                                    corrected_date = f"{corrected_date} {date_str.split(' ')[1]}"

                                parsed = datetime.strptime(corrected_date, fmt)
                                break
                        except:
                            continue

            if parsed is None:
                # 尝试pandas解析
                try:
                    import pandas as pd
                    parsed = pd.to_datetime(date_str).to_pydatetime()
                except:
                    raise ValueError(f"无法解析时间字符串: {date_str}")

            parsed_times.append(np.datetime64(parsed))

        return pr.assign_coords(time=np.array(parsed_times, dtype='datetime64[D]'))

    # ==================== 5. 未知格式 ====================
    raise ValueError(f"无法识别的时间格式，数据类型: {time_vals.dtype}")
"""======================================"""

# 检查数据格式与逐日特征
def check_data(pr):
    """
    通用降水数据校验函数（含自动时间类型转换）
    :param pr: xarray.DataArray - 待校验的降水数据
    :return: xarray.DataArray - 校验并转换后的降水数据（time为datetime64[D]格式）
    :raises ValueError: 数据校验失败时抛出异常
    """
    # 1. 基础类型校验
    if not isinstance(pr, xr.DataArray):
        raise ValueError(f"输入数据必须是xarray.DataArray类型，当前类型: {type(pr)}")

    # 2. 降水变量名校验（支持气象常用命名）
    valid_var_names = ['pr', 'pre', 'precipitation', 'rain', 'prec', 'tp']
    var_name = str(pr.name).lower() if pr.name is not None else ''
    if var_name not in valid_var_names:
        raise ValueError(
            f"无效的降水变量名: '{pr.name}'，支持的变量名: {', '.join(valid_var_names)}"
        )

    # 3. time维度存在性校验
    if 'time' not in pr.dims:
        raise ValueError("输入数据必须包含'time'维度")

    # 4. 核心：自动转换time维度为datetime64[D]格式
    pr = _convert_time_to_datetime(pr)

    # 5. 逐日数据校验（基于datetime64计算时间间隔）
    time_vals = pr.time.values
    if len(time_vals) < 2:
        raise ValueError("time维度数据量不足，至少需要2个时间点验证逐日特征")

    time_diff = np.diff(time_vals)
    day_diff = np.timedelta64(1, 'D')

    # 计算逐日比例
    daily_ratio = np.sum(time_diff == day_diff) / len(time_diff)

    # 放宽标准：允许最多5%的非逐日间隔（处理气候模型数据的不规则性）
    if daily_ratio < 0.95:
        # 手动计算时间间隔分布（兼容旧版numpy）
        unique_intervals = []
        counts = []

        for interval in time_diff:
            found = False
            for i, u_interval in enumerate(unique_intervals):
                if interval == u_interval:
                    counts[i] += 1
                    found = True
                    break
            if not found:
                unique_intervals.append(interval)
                counts.append(1)

        # 找出主要的时间间隔
        if counts:  # 确保有数据
            max_count = max(counts)
            max_index = counts.index(max_count)
            main_interval = unique_intervals[max_index]
            main_ratio = max_count / len(time_diff)

            if main_ratio >= 0.95:
                # 如果主要间隔是1天，但有一些缺失，允许继续
                if main_interval == day_diff:
                    print(f"注意: 数据有{100 * (1 - daily_ratio):.1f}%的时间间隔不是1天，但主要间隔是1天，继续处理")
                    # 不抛出异常，继续处理
                else:
                    # 主要间隔不是1天，检查是否是气候模型常用间隔
                    main_interval_days = main_interval / np.timedelta64(1, 'D')
                    if abs(main_interval_days - 1.0) < 0.01:
                        # 接近1天（如0.997天或1.003天），可能是日历差异
                        print(f"注意: 主要时间间隔为{main_interval_days:.3f}天，继续处理")
                        # 不抛出异常，继续处理
                    else:
                        # 创建间隔分布字典用于错误信息
                        interval_dist = {}
                        for interval, count in zip(unique_intervals, counts):
                            interval_days = interval / np.timedelta64(1, 'D')
                            interval_dist[f"{interval_days:.3f}天"] = count

                        raise ValueError(
                            f"输入数据非标准逐日数据！仅{daily_ratio * 100:.1f}%的时间间隔为1天\n"
                            f"主要时间间隔: {main_interval} ({main_interval_days:.3f}天)，占比{main_ratio * 100:.1f}%\n"
                            f"时间间隔分布: {interval_dist}\n"
                            f"时间间隔统计：最小={np.min(time_diff)}, 最大={np.max(time_diff)}, 均值={np.mean(time_diff)}"
                        )
            else:
                # 创建间隔分布字典用于错误信息
                interval_dist = {}
                for interval, count in zip(unique_intervals, counts):
                    interval_days = interval / np.timedelta64(1, 'D')
                    interval_dist[f"{interval_days:.3f}天"] = count

                raise ValueError(
                    f"输入数据非标准逐日数据！仅{daily_ratio * 100:.1f}%的时间间隔为1天\n"
                    f"时间间隔分布: {interval_dist}\n"
                    f"时间间隔统计：最小={np.min(time_diff)}, 最大={np.max(time_diff)}, 均值={np.mean(time_diff)}"
                )
        else:
            raise ValueError(
                f"输入数据非标准逐日数据！仅{daily_ratio * 100:.1f}%的时间间隔为1天\n"
                f"时间间隔统计：最小={np.min(time_diff)}, 最大={np.max(time_diff)}, 均值={np.mean(time_diff)}"
            )

    # 6. 年份范围有效性校验
    years = pr.time.dt.year.values
    min_year, max_year = np.min(years), np.max(years)
    if min_year > max_year:
        raise ValueError(f"无效的年份范围: 最小年份={min_year}, 最大年份={max_year}")

    return pr

"""======================================"""

# 计算连续天数辅助函数
def _max_consecutive_true(arr):
    """辅助：计算一维布尔数组中最长连续True长度（用于CDD/CWD）"""
    if arr.dtype != bool:
        arr = arr.astype(bool)

    if arr.ndim == 1:
        max_len, cur_len = 0, 0
        for val in arr:
            cur_len = cur_len + 1 if val else 0
            max_len = max(max_len, cur_len)
        return np.float32(max_len)

    # 高维数组处理（time轴为首维）
    stacked = np.moveaxis(arr, 0, -1)
    out_shape = stacked.shape[:-1]
    result = np.zeros(out_shape, dtype=np.float32)

    it = np.nditer(result, flags=['multi_index'], op_flags=[['writeonly']])
    while not it.finished:
        idx = it.multi_index
        series = stacked[idx]
        max_len, cur_len = 0, 0
        for val in series:
            cur_len = cur_len + 1 if val else 0
            max_len = max(max_len, cur_len)
        it[0] = max_len
        it.iternext()

    return result




"""=================主要函数=================================="""


# ====================== 降水变率计算函数 ======================
# ====================== 降水变率计算函数 ======================
# ====================== 降水变率计算函数 ======================
# ====================== 降水变率计算函数 ======================
# ====================== 降水变率计算函数 ======================
# ====================== 降水变率计算函数 ======================
def _calculate_cv_spatial_vectorized(data_array, include_zeros=False, min_valid_days=5):
    """
    向量化计算空间变异系数

    参数:
        data_array: numpy数组，形状为 [time, lat, lon]
        include_zeros: bool，是否包含0值数据，默认False（不考虑0值）
        min_valid_days: 最小有效数据天数要求，默认5天

    返回:
        cv_grid: numpy数组，形状为 [lat, lon]，包含每个格点的变异系数
    """
    if len(data_array.shape) != 3:
        return np.nan

    time_dim, lat_dim, lon_dim = data_array.shape

    # 在函数内部局部抑制警告
    with warnings.catch_warnings():
        # 抑制预期的numpy警告
        warnings.filterwarnings('ignore', category=RuntimeWarning,
                                message='Mean of empty slice')
        warnings.filterwarnings('ignore', category=RuntimeWarning,
                                message='Degrees of freedom <= 0 for slice')
        warnings.filterwarnings('ignore', category=RuntimeWarning,
                    message='All-NaN slice encountered')

        # 1. 根据include_zeros参数创建有效数据掩码
        if include_zeros:
            # 包含0值：只排除NaN
            valid_mask = ~np.isnan(data_array)
            calculation_type = "包含0值（干日）"
        else:
            # 不包含0值：排除NaN和0值
            valid_mask = (~np.isnan(data_array)) & (data_array > 0)
            calculation_type = "只考虑湿日（>0）"

        # 2. 计算有效数据数量
        valid_count = np.sum(valid_mask, axis=0)

        # 3. 标记有足够数据的格点
        mask_sufficient = valid_count >= min_valid_days

        # 4. 如果没有足够数据的格点，直接返回
        if not np.any(mask_sufficient):
            return np.full((lat_dim, lon_dim), np.nan)

        # 5. 提取有效数据（将无效数据设为NaN）
        valid_data = np.where(valid_mask, data_array, np.nan)

        # 6. 计算统计量
        mean_val = np.nanmean(valid_data, axis=0)
        std_val = np.nanstd(valid_data, axis=0, ddof=1)

    # 7. 计算变异系数（安全处理）
    cv_grid = np.full((lat_dim, lon_dim), np.nan)

    with np.errstate(invalid='ignore', divide='ignore'):
        cv_values = std_val / mean_val

        # 只填充满足条件的格点
        valid_cv_mask = (
                mask_sufficient &
                (valid_count > 0) &
                (mean_val > 0) &
                (~np.isnan(mean_val)) &
                (~np.isnan(std_val))
        )

        cv_grid[valid_cv_mask] = cv_values[valid_cv_mask]

        # 特殊情况处理：如果包含0值且所有数据都是0，设CV=0
        if include_zeros:
            all_zero_mask = np.all((valid_mask & (data_array == 0)) | (~valid_mask), axis=0)
            cv_grid[all_zero_mask & mask_sufficient] = 0.0

    return cv_grid


def CV_daily(pr, include_zeros=False):
    """
    计算逐年的日降水变异系数

    参数:
        pr: 降水数据（xarray DataArray）
        include_zeros: bool，是否包含0值数据，默认False（不考虑0值）

    返回:
        result: xarray DataArray，形状为 [year, lat, lon]
    """
    pr = check_data(pr)

    if include_zeros:
        print("计算逐年日降水变异系数 (CV_daily) - 包含0值数据")
        calc_desc = "包含0值（干日）"
    else:
        print("计算逐年日降水变异系数 (CV_daily) - 只考虑湿日")
        calc_desc = "只考虑湿日（>0 mm/day）"

    years = np.unique(pr.time.dt.year)
    cv_results = []
    year_labels = []

    for year in tqdm(years, desc="CV_daily"):
        year_data = pr.sel(time=pr.time.dt.year == year)

        # 检查总数据天数
        if len(year_data) < 5:
            cv_grid = np.full((len(pr.lat), len(pr.lon)), np.nan)
        else:
            # 使用向量化函数计算
            cv_grid = _calculate_cv_spatial_vectorized(
                year_data.values,
                include_zeros=include_zeros
            )

        cv_da = xr.DataArray(
            cv_grid,
            coords={'lat': pr.lat, 'lon': pr.lon},
            dims=['lat', 'lon'],
            attrs={
                'units': 'dimensionless',
                'long_name': f'Daily precipitation CV {year}',
                'description': f'Coefficient of variation (std/mean) of daily precipitation ({calc_desc})',
                'calculation': 'Excludes NaN values' + (
                    '' if include_zeros else ', excludes zero precipitation values'),
                'include_zeros': 1 if include_zeros else 0,
                'minimum_days': 5,
                'year': int(year)
            }
        )
        cv_results.append(cv_da)
        year_labels.append(int(year))

    result = xr.concat(cv_results, dim='year').assign_coords(year=year_labels)
    result.name = 'CV_daily'
    result.attrs.update({
        'calculation_type': calc_desc,
        'include_zeros': 1 if include_zeros else 0,
        'time_period': 'daily'
    })
    return result


def CV_June(pr, include_zeros=False):
    """
    计算逐年6月降水变异系数

    参数:
        pr: 降水数据
        include_zeros: bool，是否包含0值数据，默认False
    """
    pr = check_data(pr)

    if include_zeros:
        print("计算逐年6月降水变异系数 (CV_June) - 包含0值数据")
        calc_desc = "包含0值（干日）"
    else:
        print("计算逐年6月降水变异系数 (CV_June) - 只考虑湿日")
        calc_desc = "只考虑湿日（>0 mm/day）"

    years = np.unique(pr.time.dt.year)
    cv_results = []
    year_labels = []

    for year in tqdm(years, desc="CV_June"):
        # 提取6月数据
        jun_data = pr.sel(time=(pr.time.dt.year == year) & (pr.time.dt.month == 6))

        # 检查6月数据天数
        if len(jun_data) < 5:
            cv_grid = np.full((len(pr.lat), len(pr.lon)), np.nan)
        else:
            cv_grid = _calculate_cv_spatial_vectorized(
                jun_data.values,
                include_zeros=include_zeros
            )

        cv_da = xr.DataArray(
            cv_grid,
            coords={'lat': pr.lat, 'lon': pr.lon},
            dims=['lat', 'lon'],
            attrs={
                'units': 'dimensionless',
                'long_name': f'June precipitation CV {year}',
                'description': f'Coefficient of variation of June precipitation ({calc_desc})',
                'include_zeros': 1 if include_zeros else 0,
                'month': 6,
                'year': int(year)
            }
        )
        cv_results.append(cv_da)
        year_labels.append(int(year))

    result = xr.concat(cv_results, dim='year').assign_coords(year=year_labels)
    result.name = 'CV_June'
    result.attrs.update({
        'calculation_type': calc_desc,
        'include_zeros': 1 if include_zeros else 0,
        'time_period': 'monthly',
        'month': 6
    })
    return result


def CV_July(pr, include_zeros=False):
    """计算逐年7月降水变异系数"""
    pr = check_data(pr)

    if include_zeros:
        print("计算逐年7月降水变异系数 (CV_July) - 包含0值数据")
        calc_desc = "包含0值（干日）"
    else:
        print("计算逐年7月降水变异系数 (CV_July) - 只考虑湿日")
        calc_desc = "只考虑湿日（>0 mm/day）"

    years = np.unique(pr.time.dt.year)
    cv_results = []
    year_labels = []

    for year in tqdm(years, desc="CV_July"):
        jul_data = pr.sel(time=(pr.time.dt.year == year) & (pr.time.dt.month == 7))

        if len(jul_data) < 5:
            cv_grid = np.full((len(pr.lat), len(pr.lon)), np.nan)
        else:
            cv_grid = _calculate_cv_spatial_vectorized(
                jul_data.values,
                include_zeros=include_zeros
            )

        cv_da = xr.DataArray(
            cv_grid,
            coords={'lat': pr.lat, 'lon': pr.lon},
            dims=['lat', 'lon'],
            attrs={
                'units': 'dimensionless',
                'long_name': f'July precipitation CV {year}',
                'description': f'Coefficient of variation of July precipitation ({calc_desc})',
                'include_zeros': 1 if include_zeros else 0,
                'month': 7,
                'year': int(year)
            }
        )
        cv_results.append(cv_da)
        year_labels.append(int(year))

    result = xr.concat(cv_results, dim='year').assign_coords(year=year_labels)
    result.name = 'CV_July'
    result.attrs.update({
        'calculation_type': calc_desc,
        'include_zeros': 1 if include_zeros else 0,
        'time_period': 'monthly',
        'month': 7
    })
    return result


def CV_August(pr, include_zeros=False):
    """计算逐年8月降水变异系数"""
    pr = check_data(pr)

    if include_zeros:
        print("计算逐年8月降水变异系数 (CV_August) - 包含0值数据")
        calc_desc = "包含0值（干日）"
    else:
        print("计算逐年8月降水变异系数 (CV_August) - 只考虑湿日")
        calc_desc = "只考虑湿日（>0 mm/day）"

    years = np.unique(pr.time.dt.year)
    cv_results = []
    year_labels = []

    for year in tqdm(years, desc="CV_August"):
        aug_data = pr.sel(time=(pr.time.dt.year == year) & (pr.time.dt.month == 8))

        if len(aug_data) < 5:
            cv_grid = np.full((len(pr.lat), len(pr.lon)), np.nan)
        else:
            cv_grid = _calculate_cv_spatial_vectorized(
                aug_data.values,
                include_zeros=include_zeros
            )

        cv_da = xr.DataArray(
            cv_grid,
            coords={'lat': pr.lat, 'lon': pr.lon},
            dims=['lat', 'lon'],
            attrs={
                'units': 'dimensionless',
                'long_name': f'August precipitation CV {year}',
                'description': f'Coefficient of variation of August precipitation ({calc_desc})',
                'include_zeros': 1 if include_zeros else 0,
                'month': 8,
                'year': int(year)
            }
        )
        cv_results.append(cv_da)
        year_labels.append(int(year))

    result = xr.concat(cv_results, dim='year').assign_coords(year=year_labels)
    result.name = 'CV_August'
    result.attrs.update({
        'calculation_type': calc_desc,
        'include_zeros': 1 if include_zeros else 0,
        'time_period': 'monthly',
        'month': 8
    })
    return result


def CV_Summer(pr, include_zeros=False):
    """计算逐年夏季（6-8月）降水变异系数"""
    pr = check_data(pr)

    if include_zeros:
        # print("计算逐年夏季（6-8月）降水变异系数 (CV_Summer) - 包含0值数据")
        calc_desc = "包含0值（干日）"
    else:
        # print("计算逐年夏季（6-8月）降水变异系数 (CV_Summer) - 只考虑湿日")
        calc_desc = "只考虑湿日（>0 mm/day）"

    years = np.unique(pr.time.dt.year)
    cv_results = []
    year_labels = []

    for year in tqdm(years, desc="CV_Summer"):
        # 提取夏季（6-8月）数据
        summer_data = pr.sel(
            time=(pr.time.dt.year == year) & (pr.time.dt.month.isin([6, 7, 8]))
        )

        if len(summer_data) < 15:  # 夏季至少15天数据
            cv_grid = np.full((len(pr.lat), len(pr.lon)), np.nan)
        else:
            cv_grid = _calculate_cv_spatial_vectorized(
                summer_data.values,
                include_zeros=include_zeros
            )

        cv_da = xr.DataArray(
            cv_grid,
            coords={'lat': pr.lat, 'lon': pr.lon},
            dims=['lat', 'lon'],
            attrs={
                'units': 'dimensionless',
                'long_name': f'Summer precipitation CV {year}',
                'description': f'Coefficient of variation of summer (June-August) precipitation ({calc_desc})',
                'include_zeros': 1 if include_zeros else 0,
                'months': [6, 7, 8],
                'year': int(year),
                'minimum_days': 15
            }
        )
        cv_results.append(cv_da)
        year_labels.append(int(year))

    result = xr.concat(cv_results, dim='year').assign_coords(year=year_labels)
    result.name = 'CV_Summer'
    result.attrs.update({
        'calculation_type': calc_desc,
        'include_zeros': 1 if include_zeros else 0,
        'time_period': 'seasonal',
        'months': [6, 7, 8],
        'season': 'summer'
    })
    return result
