
# MI_image_numpy.py — clean, commented, NumPy‑friendly 
# -----------------------------------------------------------------------------
# This module provides a very small image processing toolkit built on NumPy.
# It is designed to be approachable for first‑year undergraduates:
#   • Images are simple Python dicts with: 'width', 'height', 'channels', 'data'
#   • Pixel data lives in a NumPy array of shape (H, W, C) with float32 values
#   • Functions operate IN‑PLACE on that dict (unless otherwise stated)
#
#
# HOW TO READ THIS FILE:
#   1) Start with the “Image container helpers” section to see how images are
#      created and copied.
#   2) Look at “Point operations” (simple per‑pixel math).
#   3) Move to “Spatial filters” (median / average / gaussian, etc.).
#   4) Explore “Frequency domain (FFT)”.
#   5) Finally, “Feature detectors” (edges, corners) build on the basics above.
# -----------------------------------------------------------------------------

from __future__ import annotations

from typing import Dict
import numpy as np
import math

# -----------------------------------------------------------------------------
# Constants and file format helpers (PNM/PPM/PGM)
# -----------------------------------------------------------------------------
PGM_RAW = 0
PGM_ASCII = 1
PPM_RAW = 2
PPM_ASCII = 3

FILE_MAGIC_TO_FORMAT = {
    'P5': PGM_RAW,   # PGM (grayscale) binary
    'P2': PGM_ASCII, # PGM (grayscale) ASCII
    'P6': PPM_RAW,   # PPM (color RGB) binary
    'P3': PPM_ASCII  # PPM (color RGB) ASCII
}

def fatal_error(message: str):
    """Stop the program with a clear error message (used for input/shape checks)."""
    print(message)
    raise SystemExit(1)

# -----------------------------------------------------------------------------
# Image container helpers
# -----------------------------------------------------------------------------
def allocate_image(width: int, height: int, channels: int, dtype=np.float32) -> Dict:
    """Create a new image dict with a zero‑initialized data array.

    Parameters
    ----------
    width, height : image geometry (must be > 0)
    channels      : number of channels (1 = gray, 3 = RGB)
    dtype         : numpy dtype for the data buffer
    """
    if width <= 0 or height <= 0 or channels not in (1, 3):
        fatal_error('allocate_image: invalid')
    data = np.zeros((height, width, channels), dtype=dtype)
    return {'width': width, 'height': height, 'channels': channels, 'data': data}

def is_empty(img: Dict) -> bool:
    """Return True if an image dict looks uninitialized/invalid."""
    return (
        img is None or
        img.get('width', 0) <= 0 or
        img.get('height', 0) <= 0 or
        img.get('channels', 0) <= 0
    )

def same_size(a: Dict, b: Dict) -> bool:
    """Check if two images share the same geometry and channel count."""
    return (
        a['width'] == b['width'] and
        a['height'] == b['height'] and
        a['channels'] == b['channels']
    )

def clone_image(img: Dict) -> Dict:
    """Deep copy an image (preserve dtype)."""
    out = allocate_image(img['width'], img['height'], img['channels'], dtype=img['data'].dtype)
    out['data'][:] = img['data']
    return out

# Convenience pixel accessors (rarely needed in vectorized code but kept for API)
def get_color(img: Dict, x: int, y: int, z: int | None = None) -> float:
    return float(img['data'][y, x, 0 if z is None else z])

def set_color(img: Dict, val: float, x: int, y: int, z: int | None = None):
    img['data'][y, x, slice(None) if z is None else z] = float(val)

def get_color_flat(img: Dict, index: int, z: int) -> float:
    w = img['width']
    y = index // w
    x = index % w
    return float(img['data'][y, x, z])

def get_interpolated_color(img: Dict, x: float, y: float, z: int) -> float:
    H, W = img['height'], img['width']
    x1 = int(np.floor(x)); y1 = int(np.floor(y))
    x2 = min(x1 + 1, W - 1); y2 = min(y1 + 1, H - 1)
    x1 = max(0, x1); y1 = max(0, y1)

    fx = x - x1; fy = y - y1
    Q11 = img['data'][y1, x1, z]
    Q21 = img['data'][y1, x2, z]
    Q12 = img['data'][y2, x1, z]
    Q22 = img['data'][y2, x2, z]

    val1 = Q11 * (1 - fx) + Q21 * fx
    val2 = Q12 * (1 - fx) + Q22 * fx
    return float(val1 * (1 - fy) + val2 * fy)

# -----------------------------------------------------------------------------
# I/O (PNM/PGM/PPM) — DO NOT MODIFY LoadImage/SaveImage (verbatim copy)
# -----------------------------------------------------------------------------
def _read_pnm_header(f):
    magic = f.readline().strip(); magic = magic.decode('ascii') if isinstance(magic,bytes) else magic
    if magic not in FILE_MAGIC_TO_FORMAT: fatal_error('LoadImage Error: Unsupported file type')
    fmt = FILE_MAGIC_TO_FORMAT[magic]
    def next_token():
        token=[]
        while True:
            ch=f.read(1)
            if not ch: break
            if ch==b'#' or ch=='#': f.readline(); continue
            if (isinstance(ch,bytes) and ch.isspace()) or (isinstance(ch,str) and ch.isspace()):
                if token: break
                else: continue
            token.append(ch)
        if not token: return None
        bts=b''.join([c if isinstance(c,bytes) else c.encode('ascii') for c in token])
        return bts.decode('ascii')
    width=int(next_token()); height=int(next_token()); maxval=int(next_token())
    return magic,fmt,width,height,maxval

def LoadImage(file_name: str) -> Dict:
    if not file_name: fatal_error('Error: Please specify a filename')
    with open(file_name,'rb') as f:
        magic,fmt,width,height,maxval=_read_pnm_header(f)
        channels=1 if magic in ('P2','P5') else 3
        img=allocate_image(width,height,channels)
        if fmt==PGM_ASCII:
            with open(file_name,'r',encoding='ascii') as fa:
                line=fa.readline(); tokens=[]
                while len(tokens)<3:
                    line=fa.readline()
                    if not line: break
                    if line.strip().startswith('#'): continue
                    tokens.extend(line.strip().split())
                values=[]
                for line in fa:
                    if line.strip().startswith('#'): continue
                    values.extend(line.strip().split())
                arr=np.array(values,dtype=np.float32)
                if arr.size<width*height: fatal_error('PGM_ASCII: unexpected end of file')
                arr=arr[:width*height].reshape((height,width))
                img['data'][:,:,0]=arr*(255.0/maxval)
                return img
        elif fmt==PGM_RAW:
            raw=np.frombuffer(f.read(width*height),dtype=np.uint8)
            if raw.size<width*height: fatal_error('PGM_RAW: unexpected end of file')
            img['data'][:,:,0]=raw.reshape((height,width))*(255.0/maxval); return img
        elif fmt==PPM_ASCII:
            with open(file_name,'r',encoding='ascii') as fa:
                line=fa.readline(); tokens=[]
                while len(tokens)<3:
                    line=fa.readline()
                    if not line: break
                    if line.strip().startswith('#'): continue
                    tokens.extend(line.strip().split())
                values=[]
                for line in fa:
                    if line.strip().startswith('#'): continue
                    values.extend(line.strip().split())
                arr=np.array(values,dtype=np.float32); expected=width*height*3
                if arr.size<expected: fatal_error('PPM_ASCII: unexpected end of file')
                arr=arr[:expected].reshape((height,width,3))
                img['data'][:,:,:]=arr*(255.0/maxval); return img
        elif fmt==PPM_RAW:
            raw=np.frombuffer(f.read(width*height*3),dtype=np.uint8)
            if raw.size<width*height*3: fatal_error('PPM_RAW: unexpected end of file')
            img['data'][:,:,:]=raw.reshape((height,width,3))*(255.0/maxval); return img
        else: fatal_error('Unsupported format')

def SaveImage(img: Dict, file_name: str, ff: int):
    if not file_name: fatal_error('Error!! SaveImage')
    H,W,C=img['data'].shape; data_u8=np.clip(img['data'],0,255).astype(np.uint8)
    if ff==PGM_ASCII:
        with open(file_name,'w',encoding='ascii') as out:
            out.write(f"P2\n{W} {H}\n255\n");
            for y in range(H): out.write(' '.join(str(int(v)) for v in data_u8[y,:,0])+'\n')
    elif ff==PGM_RAW:
        with open(file_name,'wb') as out:
            out.write(f"P5\n{W} {H}\n255\n".encode('ascii')); out.write(data_u8[:,:,0].tobytes())
    elif ff==PPM_ASCII:
        with open(file_name,'w',encoding='ascii') as out:
            out.write(f"P3\n{W} {H}\n255\n"); arr=data_u8 if C>1 else np.repeat(data_u8[:,:,0:1],3,axis=2)
            for y in range(H):
                row_vals=arr[y].reshape(-1,3); out.write(' '.join(f"{int(r)} {int(g)} {int(b)}" for r,g,b in row_vals)+'\n')
    elif ff==PPM_RAW:
        with open(file_name,'wb') as out:
            out.write(f"P6\n{W} {H}\n255\n".encode('ascii')); arr=data_u8 if C>1 else np.repeat(data_u8[:,:,0:1],3,axis=2); out.write(arr.tobytes())
    else: fatal_error('save_image: unsupported FILE_FORMAT')

def SaveSpectralImage(img: Dict, file_name: str, ff: int, logTransform: bool):
    """Save the magnitude of a complex (real/imag) spectrum as a grayscale image.

    If `logTransform` is True, we visualize with a log scale to make details
    easier to see.
    """
    H, W, C = img['data'].shape
    Re = img['data'][:, :, 0]
    Im = img['data'][:, :, 1] if C > 1 else np.zeros_like(Re)
    mag = np.sqrt(Re*Re + Im*Im) / float(W*H)
    out = np.minimum(250.0*np.log(mag + 1.0), 255.0) if logTransform else np.minimum(mag, 255.0)
    tmp = allocate_image(W, H, 1)
    tmp['data'][:, :, 0] = out
    SaveImage(tmp, file_name, ff)


# -----------------------------------------------------------------------------
# Filters
# -----------------------------------------------------------------------------
def GradientFilter(img: Dict, dir: int):
    """Compute x or y gradient using NumPy's gradient (per channel)."""

def NormGradientFilter(img: Dict, radical: bool):
    """Compute gradient magnitude (sqrt or squared magnitude)."""

def LaplacianFilter(img: Dict):
    """Apply a 3x3 Laplacian filter (edge‑padded)."""


# -----------------------------------------------------------------------------
# Feature detectors (edges, corners)
# -----------------------------------------------------------------------------
def EdgeDetec(img: Dict, threshold: float, sigma: float):
    """Simple edge detector: blur → gradient magnitude → threshold."""

def NonMaxSupp(img: Dict, sigma: float):
    """Non‑maximum suppression along gradient directions (Canny step)."""

def ZeroCrossing(img: Dict, threshold: float, sigma: float):
    """Zero‑crossing detector after Laplacian of Gaussian (LoG‑like)."""

def HarrisCornerDetec(img: Dict, halfwinsize: int, k: float):
    """Harris corner response with Gaussian window of size (2*h+1).

    Writes the (rescaled) response to the image (1 channel). For visualization.
    """

def KMeansSegmentation(img: Dict, k=2):
    """Very small 2‑cluster k‑means on a grayscale image (in‑place)."""




