import numpy as np
import math

__all__ = ["naturalize_rgb_image", "naturalize_single_channel"]

# Constants from the original code
EPS = 0.0001
T1_pr = 0.3754
N_Lap = 2041
Lap_Offset = 1020
N_Grad = 512
Grad_Offset = 256
T2_pr = [0.2421, 0.2550, 0.2474, 0.24816666]

verbose = False

def _calculate_psnr(nf):
    """Calculate PSNR from naturalization factor"""
    if 0 <= nf <= 0.934:
        return 23.65 * math.exp(0.6 * nf) - 20.0 * math.exp(-7.508 * nf)
    elif 0.934 < nf < 1.07:
        return float('inf')  # > 40
    elif 1.07 <= nf < 1.9:
        return -11.566 * nf + 52.776
    else:
        return 13.06 * nf**4 - 121.4 * nf**3 + 408.5 * nf**2 - 595.5 * nf + 349

def _find_t_error(p_t, N, offset, T):
    """Calculate error for finding optimal T parameter"""
    error = 0.0
    for i in range(-offset, N - offset):
        tmp = math.atan(T * i) - p_t[i + offset]
        error += tmp * tmp
    return error

def _find_t_golden_section(data, N, offset, eps):
    """
    Find optimal T parameter using golden section search.
    EXACT implementation matching the Java code.
    """
    left = 0.0
    right = 1.0
    
    # Create p_t array (shift and rescale) - use float32 to match Java float
    p_t = np.array([(x - 0.5) * math.pi for x in data], dtype=np.float32)
    
    m1 = 0.0
    m2 = 0.0
    
    # While the precision is bigger than eps
    while right - left >= eps:
        # Move left and right by 1/3
        m1 = left + (right - left) / 3
        m2 = right - (right - left) / 3
        
        # Evaluate on m1 and m2, and move the extreme point
        error1 = _find_t_error(p_t, N, offset, m1)
        error2 = _find_t_error(p_t, N, offset, m2)
        
        if error1 <= error2:
            right = m2
        else:
            left = m1
    
    # Return the average (this is critical - matches Java exactly)
    return (m1 + m2) / 2

def _compute_gradient_field_forward(image):
    """
    Compute forward gradient field matching Java implementation.
    G_x = I(x+1, y) - I(x, y)
    G_y = I(x, y+1) - I(x, y)
    """
    gx = np.zeros_like(image, dtype=np.float32)
    gy = np.zeros_like(image, dtype=np.float32)
    
    # Forward differences
    gx[:, :-1] = image[:, 1:] - image[:, :-1]
    gy[:-1, :] = image[1:, :] - image[:-1, :]
    
    return gx, gy

def _compute_laplacian_field(image):
    """Compute Laplacian field using 5-point stencil."""
    lap = np.zeros_like(image, dtype=np.float32)
    
    # 5-point stencil: -4*center + right + left + down + up
    lap[1:-1, 1:-1] = (
        image[1:-1, 2:] +      # right
        image[1:-1, :-2] +     # left
        image[2:, 1:-1] +      # down
        image[:-2, 1:-1] -     # up
        4 * image[1:-1, 1:-1]  # center
    )
    
    return lap

def _calculate_laplace_field_and_gradient_optimized(image, lap_cdf, grad_d):
    """
    Calculate laplacian field and gradient histogram.
    
    Key insight: Java stores grad_d[gy, gx] due to indexD[1-i]
    """
    # Ensure image is float32 in [0, 255] range
    if image.max() <= 1:
        image = (image * 255).astype(np.float32)
    else:
        image = image.astype(np.float32)

    if verbose:
        print(f"  Image stats: shape={image.shape}, min={image.min():.2f}, max={image.max():.2f}, mean={image.mean():.2f}")
    
    # Compute laplacian and gradients
    laplace_field = _compute_laplacian_field(image)
    gx, gy = _compute_gradient_field_forward(image)
    
    # Get valid regions (exclude borders)
    valid_laplace = laplace_field[1:-1, 1:-1]
    valid_gx = gx[1:-1, 1:-1]
    valid_gy = gy[1:-1, 1:-1]

    if verbose:
        print(f"  Valid region: {valid_laplace.shape}, n_pixels={valid_laplace.size}")
        print(f"  Laplacian: min={valid_laplace.min():.2f}, max={valid_laplace.max():.2f}")
        print(f"  Gradient X: min={valid_gx.min():.2f}, max={valid_gx.max():.2f}")
        print(f"  Gradient Y: min={valid_gy.min():.2f}, max={valid_gy.max():.2f}")
    
    n_pixels = valid_laplace.size
    f = np.float32(1.0 / n_pixels)  # Use float32 to match Java
    
    # Vectorized histogram binning for laplacian
    lap_bins = (valid_laplace + Lap_Offset).astype(np.int32)
    lap_bins = np.clip(lap_bins, 0, N_Lap - 1)
    np.add.at(lap_cdf, lap_bins.ravel(), f)

    if verbose:
        print(f"  Laplacian histogram: non-zero bins={np.count_nonzero(lap_cdf)}, sum={lap_cdf.sum():.6f}")
    
    # Vectorized 2D gradient histogram
    # Java: indexD[1-i] means grad_d[gy, gx]
    gx_bins = np.clip((Grad_Offset + valid_gx).astype(np.int32), 0, N_Grad - 1)
    gy_bins = np.clip((Grad_Offset + valid_gy).astype(np.int32), 0, N_Grad - 1)
    
    gx_flat = gx_bins.ravel()
    gy_flat = gy_bins.ravel()
    
    # Accumulate: grad_d[gy, gx]
    np.add.at(grad_d, (gy_flat, gx_flat), f)

    if verbose:
        print(f"  Gradient 2D histogram: non-zero bins={np.count_nonzero(grad_d)}, sum={grad_d.sum():.6f}")

def _calculate_lap_cdf(lap_cdf):
    """Convert laplacian histogram to CDF (in-place)"""
    np.cumsum(lap_cdf, out=lap_cdf, dtype=np.float32)

def _convert_grad_2d_to_cdf(grad_d):
    """
    Convert 2D gradient HISTOGRAM to 2D CDF (in-place).
    This performs cumsum in both dimensions using float32 precision.
    """
    # Ensure we're working with float32
    if grad_d.dtype != np.float32:
        grad_d = grad_d.astype(np.float32)
    
    # Row integration (cumsum along axis 0) - matches Java first loop
    np.cumsum(grad_d, axis=0, out=grad_d, dtype=np.float32)
    
    # Column integration (cumsum along axis 1) - matches Java second loop
    np.cumsum(grad_d, axis=1, out=grad_d, dtype=np.float32)

def _calculate_grad_cdf(grad_cdf, grad_d):
    """
    Calculate 1D gradient CDFs from 2D CDF.
    
    Java code integrates (sums) the 2D CDF to get marginal CDFs:
    - Gradx[gx] = sum over gy of grad_d[gy, gx]
    - Grady[gy] = sum over gx of grad_d[gy, gx]
    """
    # Integrate the 2D CDF - use float32 for consistency
    grad_x = np.sum(grad_d, axis=0, dtype=np.float32)
    grad_y = np.sum(grad_d, axis=1, dtype=np.float32)

    # Scale by 255.0 (not 256!) as per Java code
    grad_x = grad_x / np.float32(255.0)
    grad_y = grad_y / np.float32(255.0)
    
    # Store in output array
    grad_cdf[:, 0] = grad_x
    grad_cdf[:, 1] = grad_y

def _find_naturalization_factor(image, t2_prior):
    """Find optimal naturalization factor"""
    # Initialize histograms - use float32 to match Java
    lap_cdf = np.zeros(N_Lap, dtype=np.float32)
    grad_d = np.zeros((N_Grad, N_Grad), dtype=np.float32)
    
    # Calculate laplacian and gradient histograms
    _calculate_laplace_field_and_gradient_optimized(image, lap_cdf, grad_d)
    
    # Convert 2D gradient histogram to 2D CDF
    _convert_grad_2d_to_cdf(grad_d)
    
    # Extract 1D CDFs from 2D CDF
    grad_cdf = np.zeros((N_Grad, 2), dtype=np.float32)
    _calculate_grad_cdf(grad_cdf, grad_d)
    
    # Convert laplacian histogram to CDF
    _calculate_lap_cdf(lap_cdf)
    
    # Find T1 parameter - USE GOLDEN SECTION (critical fix!)
    t1_x = _find_t_golden_section(grad_cdf[:, 0], N_Grad, Grad_Offset, EPS)
    t1_y = _find_t_golden_section(grad_cdf[:, 1], N_Grad, Grad_Offset, EPS)
    t1 = (t1_x + t1_y) / (2 * T1_pr)
    
    # Find T2 parameter - USE GOLDEN SECTION (critical fix!)
    t2 = _find_t_golden_section(lap_cdf, N_Lap, Lap_Offset, EPS) / t2_prior
    
    # Calculate naturalization factor
    theta = 0.5
    nf = (1 - theta) * t1 + theta * t2

    if verbose:
        print(f"  Results: T1={t1:.6f}, T2={t2:.6f}, Nf={nf:.6f}, T2_prior={t2_prior:.4f}")
    
    return nf, t1, t2

def _naturalize_image(image, nf):
    """Apply naturalization to image"""
    # Calculate mean
    mean_original = np.mean(image)
    
    # Apply naturalization: (I - mean) * Nf + mean
    result = (image - mean_original) * nf + mean_original
    
    # Round and clamp to [0, 255] - matches Java rounding
    result = np.round(result)
    result = np.clip(result, 0, 255)
    
    return result.astype(np.uint8)

def naturalize_single_channel(image: np.ndarray, channel_idx: int = 0) -> tuple:
    """
    Naturalize a single image channel using its corresponding T2 prior.

    The T2 prior varies per channel to account for differences in the
    statistical properties of R, G, and B channels. For grayscale images,
    pass channel_idx=3 to use the averaged prior.

    Args:
        image: 2D numpy array of shape (H, W) representing a single
               channel, with pixel values in the range [0, 255].
        channel_idx: Index indicating which channel is being processed.
                     0=Red, 1=Green, 2=Blue, 3=Grayscale. Defaults to 0.

    Returns:
        tuple: A tuple of (result, nf, t1, t2, psnr) where:
            result: Naturalized channel as a uint8 numpy array of shape (H, W).
            nf:     Naturalization factor applied to this channel.
            t1:     T1 parameter derived from the gradient CDF.
            t2:     T2 parameter derived from the Laplacian CDF.
            psnr:   Estimated PSNR value based on the naturalization factor.
    """
    # Select appropriate T2 prior - matches Java indexing
    if channel_idx <= 2:
        t2_prior = T2_pr[2 - channel_idx]  # R=0->2, G=1->1, B=2->0
    else:
        t2_prior = T2_pr[3]  # Grayscale average
    
    # Find optimal naturalization factor
    nf, t1, t2 = _find_naturalization_factor(image, t2_prior)
    
    # Apply naturalization
    result = _naturalize_image(image, nf)
    
    # Calculate PSNR
    psnr = _calculate_psnr(nf)
    
    return result, nf, t1, t2, psnr

def naturalize_rgb_image(image: np.ndarray) -> tuple:
    """
    Naturalize an RGB or grayscale image by adjusting each channel independently.

    For RGB images, each channel is processed separately using its corresponding
    T2 prior. For grayscale images, a single averaged T2 prior is used instead.

    Args:
        image: Input image as a numpy array. Either (H, W) for grayscale
               or (H, W, 3) for RGB, with pixel values in the range [0, 255].

    Returns:
        tuple: A tuple of (result, nfs, t1s, t2s, psnrs) where:
            result: Naturalized image as a uint8 numpy array, same shape as input.
            nfs:    List of naturalization factors, one per channel.
            t1s:    List of T1 parameters, one per channel.
            t2s:    List of T2 parameters, one per channel.
            psnrs:  List of estimated PSNR values, one per channel.
    """
    if len(image.shape) == 2:
        # Single channel
        result, nf, t1, t2, psnr = naturalize_single_channel(image, channel_idx=3)
        return result, [nf], [t1], [t2], [psnr]
    else:
        # RGB image
        results = []
        nfs = []
        t1s = []
        t2s = []
        psnrs = []
        
        for i in range(min(3, image.shape[2])):
            channel = image[:, :, i]
            result, nf, t1, t2, psnr = naturalize_single_channel(channel, channel_idx=i)
            results.append(result)
            nfs.append(nf)
            t1s.append(t1)
            t2s.append(t2)
            psnrs.append(psnr)
        
        # Combine channels
        if len(results) == 3:
            result = np.stack(results, axis=2)
        else:
            result = results[0]
            
        return result, nfs, t1s, t2s, psnrs
