"""
Math utilities for pygndc.
"""

import math
from datetime import datetime, timedelta


def suggest_ndc_config(
    width: int,
    height: int,
    num_timestamps: int,
    base_resolution: int = 256,
    n_levels: int = 8,
    valid_pixel_ratio: float = 0.5,
    gpu_mem_gb: int = 24
):
    """
    Auto-calculate optimal NDC/K-Planes configuration parameters.
    
    Args:
        width: Image width (e.g., 8016)
        height: Image height (e.g., 4008)
        num_timestamps: Number of time steps (e.g., 365)
        base_resolution: Base resolution (default 256)
        n_levels: Number of hash levels (default 8)
        valid_pixel_ratio: Estimated valid pixel ratio (default 0.5)
        gpu_mem_gb: GPU memory size in GB (default 24)
        
    Returns:
        Dictionary with suggested parameters
    """
    print(f"\n=== NDC Configuration Auto-Calculator ===")
    print(f"Input: {width}x{height} pixels, {num_timestamps} frames.")
    print(f"Goal:  Base={base_resolution} -> Target={max(width, height)}, Levels={n_levels}")

    # Calculate per_level_scale
    target_res = max(width, height)
    if n_levels > 1:
        scale = (target_res / base_resolution) ** (1 / (n_levels - 1))
    else:
        scale = 1.0
        
    print(f"\n[1] Multi-Resolution Hierarchy:")
    print(f" -> Calculated Scale: {scale:.4f}")
    
    resolutions = []
    for i in range(n_levels):
        res = base_resolution * (scale ** i)
        resolutions.append(int(res))
    print(f" -> Resolution Steps: {resolutions}")
    
    if scale > 2.0:
        print("    (!) Warning: Scale > 2.0. MLP interpolation might be hard.")
    elif scale < 1.4:
        print("    (!) Warning: Scale < 1.4. Levels are too dense.")
    else:
        print("    (OK) Scale is in the golden range [1.4 - 2.0].")

    # Calculate xy_hash_size
    total_pixels = width * height
    estimated_valid_pixels = total_pixels * valid_pixel_ratio
    
    ideal_log2_xy = math.log2(estimated_valid_pixels)
    suggested_xy = int(math.ceil(ideal_log2_xy))
    
    max_xy_limit = 24 if gpu_mem_gb >= 24 else 21
    if gpu_mem_gb >= 40:
        max_xy_limit = 26
    
    final_xy = min(suggested_xy, max_xy_limit)
    final_xy = max(final_xy, 19)

    print(f"\n[2] Spatial Hashmap (XY Plane):")
    print(f" -> Estimated Valid Pixels: {int(estimated_valid_pixels):,}")
    print(f" -> Ideal Log2 Size: {ideal_log2_xy:.2f}")
    print(f" -> Suggested xy_hash_size: {final_xy} (Capacity: {2**final_xy:,})")

    # Calculate t_hash_size
    xt_entries = width * num_timestamps * valid_pixel_ratio
    yt_entries = height * num_timestamps * valid_pixel_ratio
    max_t_entries = max(xt_entries, yt_entries)
    
    ideal_log2_t = math.log2(max_t_entries)
    suggested_t = int(math.ceil(ideal_log2_t))
    
    final_t = min(max(suggested_t, 17), 22)
    
    print(f"\n[3] Temporal Hashmap (XT/YT Planes):")
    print(f" -> Estimated Spatio-Temporal Entries: {int(max_t_entries):,}")
    print(f" -> Ideal Log2 Size: {ideal_log2_t:.2f}")
    print(f" -> Suggested t_hash_size: {final_t} (Capacity: {2**final_t:,})")

    config = {
        "n_levels": n_levels,
        "base_resolution": base_resolution,
        "per_level_scale": round(scale, 4),
        "xy_hash_size": final_xy,
        "t_hash_size": final_t,
        "xt_hash_size": final_t,
        "yt_hash_size": final_t,
    }
    
    print(f"\n=== Final Recommended Config ===")
    print(config)
    return config


def calculate_time_params(start_date_str, base_res, n_levels, scale, days_per_cell=5.0, smooth_interval_days=1.0):
    """
    Calculate global end date and smooth epsilon.
    
    Args:
        start_date_str: Start date 'YYYY-MM-DD'
        base_res: K-Planes base resolution (e.g., 16)
        n_levels: Number of levels (e.g., 4)
        scale: Level scale factor (e.g., 7)
        days_per_cell: Days per cell at highest resolution (default 5.0)
        smooth_interval_days: Physical time interval for smoothness (default 1.0)
        
    Returns:
        Tuple of (end_date_str, smooth_eps)
    """
    n_max = int(base_res * (scale ** (n_levels - 1)))
    required_total_days = n_max * days_per_cell
    
    start_date = datetime.strptime(start_date_str, "%Y-%m-%d")
    end_date = start_date + timedelta(days=required_total_days)
    
    if required_total_days > 0:
        smooth_eps = smooth_interval_days / required_total_days
    else:
        smooth_eps = 0.001
    
    print(f"--- GeoNDC Time Parameters ---")
    print(f"Grid Max Cells:      {n_max}")
    print(f"Resolution:          {days_per_cell} days/cell")
    print(f"Total Span:          {required_total_days:.1f} days")
    print(f"Global Start:        {start_date_str}")
    print(f"Global End:          {end_date.strftime('%Y-%m-%d')}")
    print(f"Smooth Interval:     {smooth_interval_days} days (Physical)")
    print(f"Calculated Eps:      {smooth_eps:.8f} (Normalized)")
    print(f"------------------------------")
    
    return end_date.strftime('%Y-%m-%d'), smooth_eps

if __name__ == "__main__":
    # Example usage
    calculate_time_params(
        start_date_str="2010-02-01",
        base_res=16,
        n_levels=11,
        scale=1.41,
        days_per_cell=10.0,
        smooth_interval_days=1.0
    )
