# -*- coding: utf-8 -*-
"""
Forest Inventory Plot Designer — Sample Size Calculator.

Implements sampling formulas from:
  - Thông tư 05 / 2024 / TT-BNNPTNT (Định mức KTKT Điều tra rừng)
  - QĐ 243 — Sổ tay KTNT Kiểm kê rừng
  - HDKT ĐTR — Hướng dẫn kỹ thuật điều tra rừng
  - GIZ-SFM Vietnam 2024 (Vũ Tấn Phương)
  - Winrock International SOP 2018

Author: Lộc Vũ Trung (LVT) / Slow Forest
License: GPL-3.0
"""

import math


# ─── Default Parameters ─────────────────────────────────────────────

T_SQUARED = 4            # t² at 95% confidence
DEFAULT_ERROR_PCT = 10   # Δ% — default allowable error
DEFAULT_CV_PCT = 30      # S% — default coefficient of variation
PLOT_AREA_TIMBER = 1000  # m² — concentric circular for timber
PLOT_AREA_BAMBOO_SCATTERED = 100   # m² — bamboo mọc tản
PLOT_AREA_BAMBOO_CLUMP = 1000      # m² — bamboo mọc cụm
PLOT_AREA_GROWTH = 100   # m² — TT33 / 2018 growth survey (R=5.64m)
TT33_SAMPLING_RATE = 0.0001  # 0.01% per age class
PLOTS_PER_DAY = 4        # ô/ngày / nhóm
CREW_SIZE = 4            # người / tổ
NUM_CREWS = 2            # 2 tổ mặc định


class StratumConfig:
    """Configuration for a single forest stratum."""

    def __init__(self, name, code, area_ha, cv_pct=DEFAULT_CV_PCT,
                 forest_type='timber', plot_area_m2=None):
        """
        Args:
            name: Display name (e.g. "Rừng lá rộng TX giàu").
            code: Short code (e.g. "TXG").
            area_ha: Area in hectares.
            cv_pct: Coefficient of variation (S%), default 30%.
            forest_type: 'timber', 'bamboo_scattered', 'bamboo_clump',
                         'plantation', 'growth_plantation'.
            plot_area_m2: Override plot area
            if None auto-determined.
        """
        self.name = name
        self.code = code
        self.area_ha = float(area_ha) if area_ha else 0.0
        self.cv_pct = float(cv_pct) if cv_pct else DEFAULT_CV_PCT
        self.forest_type = forest_type

        if plot_area_m2 is not None:
            self.plot_area_m2 = float(plot_area_m2)
        elif forest_type == 'bamboo_scattered':
            self.plot_area_m2 = PLOT_AREA_BAMBOO_SCATTERED
        elif forest_type == 'bamboo_clump':
            self.plot_area_m2 = PLOT_AREA_BAMBOO_CLUMP
        elif forest_type == 'growth_plantation':
            self.plot_area_m2 = PLOT_AREA_GROWTH
        else:
            self.plot_area_m2 = PLOT_AREA_TIMBER


def compute_sample_size(stratum, t_sq=T_SQUARED, error_pct=DEFAULT_ERROR_PCT):
    """Calculate the required number of plots for one stratum.

    Implements TT05 / 2024 / TT-BNNPTNT:

    a) Random sampling (DT ≥ 2000 ha for natural forest):
        Ni = t² × (S%i)² / (Δ%)²

    b) Typical sampling (DT < 2000 ha for natural forest):
        < 150 ha  → min 30 ODD
        150–200   → 0.2% sampling rate
        > 200     → rate decreases 0.005% per 100 ha increment

    c) Plantation forest:
        ≤ 100 ha  → 5 ODD per species per age class
        > 100 ha  → 10 ODD per species per age class

    Args:
        stratum: A StratumConfig instance.
        t_sq: t-squared value (default 4 for 95% CI).
        error_pct: Allowable error percentage (default 10%).

    Returns:
        dict with method, n_plots, sampling_rate_pct, survey_area_ha,
        formula_display.
    """
    area = stratum.area_ha
    plot_area_ha = stratum.plot_area_m2 / 10_000

    if area <= 0:
        return {
            'method': 'n / a',
            'n_plots': 0,
            'sampling_rate_pct': 0.0,
            'survey_area_ha': 0.0,
            'formula_display': '—',
        }

    # Growth survey — TT33/2018 §Điều 11 Khoản 2c
    # Rừng trồng thuần loài HGĐ: tỷ lệ rút mẫu > 0.01% / cấp tuổi
    # Ô tròn 100 m² (R=5.64m), đo toàn bộ Hvn + D1.3
    if stratum.forest_type == 'growth_plantation':
        n_raw = area * TT33_SAMPLING_RATE / plot_area_ha
        n_plots = max(int(math.ceil(n_raw)), 1)
        survey_area = n_plots * plot_area_ha
        rate = (survey_area / area) * 100 if area > 0 else 0
        return {
            'method': 'growth',
            'n_plots': n_plots,
            'sampling_rate_pct': round(rate, 4),
            'survey_area_ha': round(survey_area, 2),
            'formula_display': (
                f'TT33: {area} ha × 0.01% / {plot_area_ha} ha '
                f'= {n_raw:.1f} → {n_plots} ODD / cấp tuổi'
            ),
        }

    # Plantation forest — special rules per TT05 §1.3.1(b)
    if stratum.forest_type == 'plantation':
        if area <= 100:
            n_plots = 5
            formula = 'Rừng trồng ≤ 100 ha → 5 ODD / loài / tuổi'
        else:
            n_plots = 10
            formula = 'Rừng trồng > 100 ha → 10 ODD / loài / tuổi'
        survey_area = n_plots * plot_area_ha
        rate = (survey_area / area) * 100 if area > 0 else 0
        return {
            'method': 'plantation',
            'n_plots': n_plots,
            'sampling_rate_pct': round(rate, 4),
            'survey_area_ha': round(survey_area, 2),
            'formula_display': formula,
        }

    # Natural forest — random sampling (≥ 2000 ha)
    if area >= 2000:
        cv = stratum.cv_pct
        n = t_sq * (cv ** 2) / (error_pct ** 2)
        n_plots = max(int(math.ceil(n)), 30)
        survey_area = n_plots * plot_area_ha
        rate = (survey_area / area) * 100 if area > 0 else 0

        return {
            'method': 'random',
            'n_plots': n_plots,
            'sampling_rate_pct': round(rate, 4),
            'survey_area_ha': round(survey_area, 2),
            'formula_display': (
                f'Ni = {t_sq} × {cv}² / {error_pct}² '
                f'= {n:.1f} → {n_plots}'
            ),
        }

    # Natural forest — typical sampling (< 2000 ha)
    if area < 150:
        n_plots = 30  # TT05: min 30 ODD
        formula = 'DT < 150 ha → min 30 ODD (TT05)'
    elif area <= 200:
        rate = 0.2
        survey_area_raw = area * rate / 100
        n_plots = max(int(math.ceil(survey_area_raw / plot_area_ha)), 30)
        formula = '150 ≤ DT ≤ 200 → tỷ lệ 0.2%'
    else:
        # > 200 ha: rate decreases 0.005% per 100 ha increment
        increments = math.ceil((area - 200) / 100)
        rate = max(0.2 - increments * 0.005, 0.05)
        survey_area_raw = area * rate / 100
        n_plots = max(int(math.ceil(survey_area_raw / plot_area_ha)), 30)
        formula = (
            f'DT > 200 ha → 0.2% − {increments}×0.005% '
            f'= {rate:.3f}%'
        )

    survey_area = n_plots * plot_area_ha
    effective_rate = (survey_area / area) * 100 if area > 0 else 0

    return {
        'method': 'typical',
        'n_plots': n_plots,
        'sampling_rate_pct': round(effective_rate, 4),
        'survey_area_ha': round(survey_area, 2),
        'formula_display': formula,
    }


def compute_all_strata(strata_list, t_sq=T_SQUARED,
                       error_pct=DEFAULT_ERROR_PCT):
    """Calculate sample sizes for all strata.

    Args:
        strata_list: list of StratumConfig.
        t_sq: t-squared value.
        error_pct: allowable error percentage.

    Returns:
        list of dicts (one per stratum) with extra keys:
            stratum_name, stratum_code, area_ha
        Also appended: a summary dict with totals.
    """
    results = []
    total_area = 0.0
    total_plots = 0
    total_survey = 0.0

    for s in strata_list:
        r = compute_sample_size(s, t_sq, error_pct)
        r['stratum_name'] = s.name
        r['stratum_code'] = s.code
        r['area_ha'] = s.area_ha
        r['cv_pct'] = s.cv_pct
        r['plot_area_m2'] = s.plot_area_m2
        r['forest_type'] = s.forest_type
        results.append(r)
        total_area += s.area_ha
        total_plots += r['n_plots']
        total_survey += r['survey_area_ha']

    # Summary
    overall_rate = (total_survey / total_area * 100) if total_area > 0 else 0
    work_days = math.ceil(total_plots / (PLOTS_PER_DAY * NUM_CREWS))

    summary = {
        'total_area_ha': round(total_area, 2),
        'total_plots': total_plots,
        'total_survey_ha': round(total_survey, 2),
        'overall_rate_pct': round(overall_rate, 4),
        'est_work_days': work_days,
        'crews': NUM_CREWS,
        'plots_per_day': PLOTS_PER_DAY,
        'crew_size': CREW_SIZE,
    }
    return results, summary


# ─── Slope Correction ───────────────────────────────────────────────

def slope_corrected_radius(nest_radius_m, slope_degrees):
    """Calculate slope-corrected plot radius (Winrock SOP).

    When slope > 10%, the field radius must be expanded so that
    the horizontal projection equals the designed plot area.

    Sloped_Radius = Nest_Radius / cos(slope)
    Corrected_Area = π × Nest_Radius × Sloped_Radius
    Corrected_Radius = √(Corrected_Area / π)

    Args:
        nest_radius_m: Designed radius on flat terrain (m).
        slope_degrees: Slope angle in degrees.

    Returns:
        dict with sloped_radius, corrected_area, corrected_radius.
    """
    slope_pct = math.tan(math.radians(slope_degrees)) * 100
    if slope_pct <= 10:
        # No correction needed
        area = math.pi * nest_radius_m ** 2
        return {
            'needs_correction': False,
            'sloped_radius': nest_radius_m,
            'corrected_area': round(area, 2),
            'corrected_radius': nest_radius_m,
            'slope_pct': round(slope_pct, 1),
        }

    cos_slope = math.cos(math.radians(slope_degrees))
    sloped_r = nest_radius_m / cos_slope
    corrected_area = math.pi * nest_radius_m * sloped_r
    corrected_r = math.sqrt(corrected_area / math.pi)

    return {
        'needs_correction': True,
        'sloped_radius': round(sloped_r, 2),
        'corrected_area': round(corrected_area, 2),
        'corrected_radius': round(corrected_r, 2),
        'slope_pct': round(slope_pct, 1),
    }


# ─── Plot Dimension Presets ──────────────────────────────────────────

# GIZ-SFM Vietnam concentric circular plot (3-nest)
GIZ_CONCENTRIC = [
    {'label': 'Ô 100 m² (DBH 6–22 cm)', 'radius_m': 5.64, 'area_m2': 100,
     'dbh_min': 6, 'dbh_max': 22},
    {'label': 'Ô 500 m² (DBH 22–42 cm)', 'radius_m': 12.62, 'area_m2': 500,
     'dbh_min': 22, 'dbh_max': 42},
    {'label': 'Ô 1000 m² (DBH ≥ 42 cm)', 'radius_m': 17.84, 'area_m2': 1000,
     'dbh_min': 42, 'dbh_max': None},
]

# Winrock International circular nested plot (4-nest)
WINROCK_CIRCULAR = [
    {'label': 'Sapling (R=2m)', 'radius_m': 2, 'area_m2': 12.57,
     'dbh_min': 0, 'dbh_max': 5},
    {'label': 'Small trees (R=6m)', 'radius_m': 6, 'area_m2': 113.1,
     'dbh_min': 5, 'dbh_max': 20},
    {'label': 'Medium trees (R=14m)', 'radius_m': 14, 'area_m2': 615.75,
     'dbh_min': 20, 'dbh_max': 50},
    {'label': 'Large trees (R=20m)', 'radius_m': 20, 'area_m2': 1256.64,
     'dbh_min': 50, 'dbh_max': None},
]

# Winrock International nested square plot (3-nest + sapling)
WINROCK_SQUARE = [
    {'label': 'Sapling (3×3m)', 'side_m': 3, 'area_m2': 9,
     'dbh_min': 0, 'dbh_max': 5},
    {'label': 'Small (7×7m, DBH 5–20)', 'side_m': 7, 'area_m2': 49,
     'dbh_min': 5, 'dbh_max': 20},
    {'label': 'Medium (25×25m, DBH 20–50)', 'side_m': 25, 'area_m2': 625,
     'dbh_min': 20, 'dbh_max': 50},
    {'label': 'Large (35×35m, DBH >50)', 'side_m': 35, 'area_m2': 1225,
     'dbh_min': 50, 'dbh_max': None},
]

# TT05/2024 standard circular plot (single ring)
TT05_CIRCULAR = [
    {'label': 'ODD 1000 m² (D1.3 ≥ 6 cm)', 'radius_m': 17.84,
     'area_m2': 1000, 'dbh_min': 6, 'dbh_max': None},
]

# TT05 mangrove/swamp plot
TT05_MANGROVE = [
    {'label': 'ODD 100 m² (ngập mặn / phèn)', 'radius_m': 5.64,
     'area_m2': 100, 'dbh_min': 4, 'dbh_max': None},
]

# TT05 bamboo scattered plot
TT05_BAMBOO_SCATTERED = [
    {'label': 'Ô phụ 100 m² (tre nứa tản)', 'radius_m': 5.64,
     'area_m2': 100, 'dbh_min': 2, 'dbh_max': None},
]

# Household plantation plots — circular (multiple sizes)
HOUSEHOLD_CIRCULAR = [
    {'label': 'Ô tròn 100 m² (R=5.64 m)', 'radius_m': 5.64,
     'area_m2': 100, 'dbh_min': 6, 'dbh_max': None},
    {'label': 'Ô tròn 200 m² (R=7.98 m)', 'radius_m': 7.98,
     'area_m2': 200, 'dbh_min': 6, 'dbh_max': None},
    {'label': 'Ô tròn 400 m² (R=11.28 m)', 'radius_m': 11.28,
     'area_m2': 400, 'dbh_min': 6, 'dbh_max': None},
    {'label': 'Ô tròn 500 m² (R=12.62 m)', 'radius_m': 12.62,
     'area_m2': 500, 'dbh_min': 6, 'dbh_max': None},
]

# Household plantation plots — square (multiple sizes)
HOUSEHOLD_SQUARE = [
    {'label': 'Ô vuông 100 m² (10×10 m)', 'side_m': 10,
     'area_m2': 100, 'dbh_min': 6, 'dbh_max': None},
    {'label': 'Ô vuông 200 m² (14.14×14.14 m)', 'side_m': 14.14,
     'area_m2': 200, 'dbh_min': 6, 'dbh_max': None},
    {'label': 'Ô vuông 400 m² (20×20 m)', 'side_m': 20,
     'area_m2': 400, 'dbh_min': 6, 'dbh_max': None},
    {'label': 'Ô vuông 500 m² (22.36×22.36 m)', 'side_m': 22.36,
     'area_m2': 500, 'dbh_min': 6, 'dbh_max': None},
]

PLOT_PRESETS = {
    'tt05_circular': {
        'name': 'Ô tròn TT05 / 2024 (1000 m²)',
        'name_en': 'TT05 / 2024 Circular (1000 m²)',
        'shape': 'circle',
        'nests': TT05_CIRCULAR,
        'legal_ref': 'TT 05 / 2024 / TT-BNNPTNT',
    },
    'tt33_growth': {
        'name': 'ĐT tăng trưởng HGĐ — TT33 (100 m²)',
        'name_en': 'Household Growth Survey — TT33 (100 m²)',
        'shape': 'circle',
        'nests': [
            {'label': 'OTC 100 m² (R=5.64m, đo toàn bộ Hvn+D1.3)',
             'radius_m': 5.64, 'area_m2': 100,
             'dbh_min': 0, 'dbh_max': None,
             'measures': ['D1.3', 'Hvn', 'phẩm chất']},
        ],
        'legal_ref': 'TT 33 / 2018 / TT-BNNPTNT, Điều 11 Khoản 2c',
        'note': 'Tỷ lệ rút mẫu > 0.01%/cấp tuổi, từ tuổi 3+',
    },
    'tt05_mangrove': {
        'name': 'Ô tròn rừng ngập mặn (100 m²)',
        'name_en': 'Mangrove Plot (100 m²)',
        'shape': 'circle',
        'nests': TT05_MANGROVE,
        'legal_ref': 'TT 05 / 2024 / TT-BNNPTNT',
    },
    'giz_concentric': {
        'name': 'Ô tròn đồng tâm (GIZ-SFM VN)',
        'name_en': 'Concentric Circular (GIZ-SFM VN)',
        'shape': 'circle',
        'nests': GIZ_CONCENTRIC,
    },
    'winrock_circular': {
        'name': 'Ô tròn Winrock (4 lớp)',
        'name_en': 'Winrock Circular (4-nest)',
        'shape': 'circle',
        'nests': WINROCK_CIRCULAR,
    },
    'winrock_square': {
        'name': 'Ô vuông lồng nhau Winrock',
        'name_en': 'Winrock Nested Square',
        'shape': 'square',
        'nests': WINROCK_SQUARE,
    },
    'household_circular': {
        'name': 'Rừng trồng HGĐ — Ô tròn (100–500 m²)',
        'name_en': 'Household Plantation — Circular (100–500 m²)',
        'shape': 'circle',
        'nests': HOUSEHOLD_CIRCULAR,
    },
    'household_square': {
        'name': 'Rừng trồng HGĐ — Ô vuông (100–500 m²)',
        'name_en': 'Household Plantation — Square (100–500 m²)',
        'shape': 'square',
        'nests': HOUSEHOLD_SQUARE,
    },
}
