"""
Edge-aware Poisson image editing with a Perona-Malik diffusion operator that
inhibits blending across strong edges in either source or target. Two modes:
  - pie: classical Poisson editing 
  - ours: edge-aware diffusion from source/target gradients. The parameter kappa controls the strength of the diffusion.

Usage:
    python main.py --data_dir data/cloud --method ours --kappa 0.05
"""

import os
import argparse
import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as spla
from dataclasses import dataclass
from typing import Tuple
from PIL import Image

import utils


# =============================================================================
# Configuration
# =============================================================================

@dataclass
class EdgeAwareConfig:
    """Configuration for edge-aware Poisson editing methods."""
    kappa: float = 0.02              # Perona-Malik contrast parameter
    gradient_mode: str = "max"       # Mixed gradient rule: "max", "alpha", or "source"
    gradient_alpha: float = 1.0      # Alpha for gradient mixing when mode="alpha"


# =============================================================================
# Helper Functions
# =============================================================================

def _ensure_3d(img: np.ndarray) -> Tuple[np.ndarray, bool]:
    """Return array as float64 with channel axis first (C, H, W)."""
    img = np.asarray(img, dtype=np.float64)
    was_gray = img.ndim == 2
    if was_gray:
        img = img[..., None]
    channel_first = np.moveaxis(img, -1, 0)
    return channel_first, was_gray


def _restore_channels(channels: np.ndarray, was_gray: bool) -> np.ndarray:
    """Restore array from (C, H, W) to (H, W, C) or (H, W) for grayscale."""
    img = np.moveaxis(channels, 0, -1)
    if was_gray:
        return img[..., 0]
    return img


def _prepare_mask(mask: np.ndarray) -> np.ndarray:
    """Validate and convert mask to boolean array."""
    mask = np.asarray(mask).astype(bool)
    if mask.ndim != 2:
        raise ValueError("Mask must be a 2D array")
    if mask.sum() == 0:
        raise ValueError("Mask contains no active pixels")
    return mask


# =============================================================================
# Linear System Construction
# =============================================================================

def solve_linear_system(A: sp.csr_matrix, b: np.ndarray) -> np.ndarray:
    """Solve sparse linear system Ax = b."""
    try:
        print("Solving linear system...")
        return spla.spsolve(A, b)
    except Exception as e:
        print(f"Solver failed: {e}")
        return np.zeros_like(b)


def build_weighted_laplacian(mask: np.ndarray, wx: np.ndarray, wy: np.ndarray) -> sp.csr_matrix:
    """Assemble sparse Laplacian matrix using edge weights wx/wy."""
    mask_bool = mask.astype(bool)
    idx = utils.build_mask_index(mask_bool.astype(np.int32))
    h, w = mask_bool.shape
    rows, cols, data = [], [], []

    for y in range(h):
        for x in range(w):
            if not mask_bool[y, x]:
                continue
            row = idx[y, x]
            diag = 0.0

            # Left neighbor
            if x > 0:
                weight = wx[y, x - 1]
                diag += weight
                if mask_bool[y, x - 1]:
                    rows.append(row)
                    cols.append(idx[y, x - 1])
                    data.append(-weight)
            # Right neighbor
            if x < w - 1:
                weight = wx[y, x]
                diag += weight
                if mask_bool[y, x + 1]:
                    rows.append(row)
                    cols.append(idx[y, x + 1])
                    data.append(-weight)
            # Top neighbor
            if y > 0:
                weight = wy[y - 1, x]
                diag += weight
                if mask_bool[y - 1, x]:
                    rows.append(row)
                    cols.append(idx[y - 1, x])
                    data.append(-weight)
            # Bottom neighbor
            if y < h - 1:
                weight = wy[y, x]
                diag += weight
                if mask_bool[y + 1, x]:
                    rows.append(row)
                    cols.append(idx[y + 1, x])
                    data.append(-weight)

            rows.append(row)
            cols.append(row)
            data.append(diag)

    n = mask_bool.sum()
    return sp.csr_matrix((data, (rows, cols)), shape=(n, n))


def build_divergence_rhs(vx: np.ndarray, vy: np.ndarray, 
                         target: np.ndarray, mask: np.ndarray, 
                         wx: np.ndarray, wy: np.ndarray) -> np.ndarray:
    """
    Build RHS vector by calculating weighted divergence of vector field (vx, vy)
    and adding Dirichlet boundary conditions from target image.
    """
    h, w = mask.shape
    b = []

    for y in range(h):
        for x in range(w):
            if not mask[y, x]:
                continue
            
            divergence_val = 0.0
            boundary_val = 0.0
            
            # Left neighbor
            if x > 0:
                weight = wx[y, x - 1]
                divergence_val += weight * vx[y, x - 1]
                if not mask[y, x - 1]:
                    boundary_val += weight * target[y, x - 1]

            # Right neighbor
            if x < w - 1:
                weight = wx[y, x]
                divergence_val -= weight * vx[y, x]
                if not mask[y, x + 1]:
                    boundary_val += weight * target[y, x + 1]

            # Top neighbor
            if y > 0:
                weight = wy[y - 1, x]
                divergence_val += weight * vy[y - 1, x]
                if not mask[y - 1, x]:
                    boundary_val += weight * target[y - 1, x]

            # Bottom neighbor
            if y < h - 1:
                weight = wy[y, x]
                divergence_val -= weight * vy[y, x]
                if not mask[y + 1, x]:
                    boundary_val += weight * target[y + 1, x]

            b.append(divergence_val + boundary_val)

    return np.array(b, dtype=np.float64)


def compute_mixed_gradients(source, target, mode="alpha", alpha=1.0):
    """
    Compute mixed gradient field from source and target images.
    
    Args:
        source: Source image
        target: Target image
        mode: "max" (keep stronger gradient) or "alpha" (blend gradients)
        alpha: Blending weight when mode="alpha"
    
    
    Returns:
        vx, vy: Horizontal and vertical gradient fields
    """
    h, w = source.shape
    
    # Forward differences
    s_dx = np.zeros_like(source)
    s_dx[:, :-1] = source[:, 1:] - source[:, :-1]
    t_dx = np.zeros_like(target)
    t_dx[:, :-1] = target[:, 1:] - target[:, :-1]
    
    s_dy = np.zeros_like(source)
    s_dy[:-1, :] = source[1:, :] - source[:-1, :]
    t_dy = np.zeros_like(target)
    t_dy[:-1, :] = target[1:, :] - target[:-1, :]

    if mode == "max":
        # Keep whichever gradient has larger magnitude
        mask_x = np.abs(t_dx) > np.abs(s_dx)
        mask_y = np.abs(t_dy) > np.abs(s_dy)
        vx = np.where(mask_x, t_dx, s_dx)
        vy = np.where(mask_y, t_dy, s_dy)
    elif mode == "alpha":
        vx = alpha * s_dx + (1 - alpha) * t_dx
        vy = alpha * s_dy + (1 - alpha) * t_dy
    else:
        raise ValueError(f"Unsupported gradient mode '{mode}'")
    
    return vx, vy


# =============================================================================
# Poisson Editing Methods
# =============================================================================

def run_pie(src_rgb, target_rgb, mask, config: EdgeAwareConfig):
    """
    Classical Poisson editing with unit conductivities.
    This is the baseline method without edge-awareness.
    """
    mask = _prepare_mask(mask)
    src_channels, was_gray = _ensure_3d(src_rgb)
    target_channels, _ = _ensure_3d(target_rgb)
    
    # Unit weights (no edge-awareness)
    wx = np.ones((mask.shape[0], mask.shape[1] - 1), dtype=np.float64)
    wy = np.ones((mask.shape[0] - 1, mask.shape[1]), dtype=np.float64)
    
    A = build_weighted_laplacian(mask, wx, wy)

    solved = []
    for src_channel, tgt_channel in zip(src_channels, target_channels):
        vx, vy = compute_mixed_gradients(src_channel, tgt_channel, config.gradient_mode, config.gradient_alpha)
        rhs = build_divergence_rhs(vx, vy, tgt_channel, mask, wx, wy)
        x = solve_linear_system(A, rhs)
        
        sol_channel = np.array(tgt_channel, copy=True)
        sol_channel[mask] = x
        solved.append(sol_channel)
    
    solved = np.stack(solved, axis=0)
    return np.clip(_restore_channels(solved, was_gray), 0.0, 1.0)


def run_ours(src_rgb, target_rgb, mask, config: EdgeAwareConfig):
    """
    Fixed guidance, mixed-gradient weighted Poisson.
    Uses Perona-Malik conductivity derived from source/target luminance.
    """
    mask = _prepare_mask(mask)
    src_channels, was_gray = _ensure_3d(src_rgb)
    target_channels, _ = _ensure_3d(target_rgb)

    # Compute conductivities from both images, take minimum
    src_lum = utils.rgb2lum(src_rgb)
    src_conductivity = utils.perona_malik_conductivity_from_guidance(src_lum, config.kappa)
    
    target_lum = utils.rgb2lum(target_rgb)
    target_conductivity = utils.perona_malik_conductivity_from_guidance(target_lum, config.kappa)
    
    conductivity = np.minimum(src_conductivity, target_conductivity)
    wx, wy = utils.edge_weights_from_conductivity(conductivity)
    
    A = build_weighted_laplacian(mask, wx, wy)

    solved = []
    for src_channel, tgt_channel in zip(src_channels, target_channels):
        vx, vy = compute_mixed_gradients(src_channel, tgt_channel, config.gradient_mode, config.gradient_alpha)
        rhs = build_divergence_rhs(vx, vy, tgt_channel, mask, wx, wy)
        x = solve_linear_system(A, rhs)
        sol_channel = np.array(tgt_channel, copy=True)
        sol_channel[mask] = x
        solved.append(sol_channel)
    
    solved = np.stack(solved, axis=0)
    return np.clip(_restore_channels(solved, was_gray), 0.0, 1.0)


# =============================================================================
# CLI Interface
# =============================================================================

def parse_args():
    parser = argparse.ArgumentParser(
        description="Edge-aware Poisson / Perona-Malik image editing.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--data_dir", type=str, required=True, 
                        help="Folder containing mask/source/target images.")
    parser.add_argument("--method", type=str, required=True, choices=["pie", "ours"], 
                        help="Method: pie=classical Poisson, ours=edge-aware PM weighted")
    parser.add_argument("--kappa", type=float, default=0.02, 
                        help="Perona-Malik contrast parameter κ.")
    parser.add_argument("--gradient-mode", type=str, default="alpha", choices=["max", "alpha", "source"], 
                        help="Mixed gradient rule.")
    parser.add_argument("--gradient-alpha", type=float, default=1.0, 
                        help="Alpha for gradient mixing when mode='alpha'.")
    parser.add_argument("--output", type=str, 
                        help="Optional output filename.")
    return parser.parse_args()


def load_inputs(args):
    """Load source, target, and mask images from data directory."""
    mask = utils.read_image(args.data_dir, "mask", gray=True)
    src = utils.read_image(args.data_dir, "source", gray=False)
    tgt = utils.read_image(args.data_dir, "target", gray=False)
    mask = (mask > 0.5).astype(np.uint8)
    return src, tgt, mask


def save_output(image, path):
    """Save result image to disk."""
    image = np.clip(image, 0.0, 1.0)
    array = (image * 255).astype(np.uint8)
    Image.fromarray(array).save(path)


def main():
    try:
        args = parse_args()
        src, tgt, mask = load_inputs(args)

        config = EdgeAwareConfig(
            kappa=args.kappa,
            gradient_mode=args.gradient_mode,
            gradient_alpha=args.gradient_alpha,
        )

        methods = {
            "pie": run_pie,
            "ours": run_ours,
        }

        result = methods[args.method](src, tgt, mask, config)

        # Save output with descriptive default names
        if args.method == "ours":
            output_name = args.output or f"result_ours_kappa_{args.kappa}.png"
        else:
            output_name = args.output or "result_pie.png"

        output_path = output_name if os.path.isabs(output_name) else os.path.join(args.data_dir, output_name)
        save_output(result, output_path)
        print(f"Saved result to {output_path}")

    except Exception as e:
        print(f"Error: {e}")
        import sys
        sys.exit(0)


if __name__ == "__main__":
    main()

