"""
Total Area Coverage (TAC) and CMYK Channel Separation Visualizer
Generates previews, channel isolations, and coverage heatmap overlays.
"""

import io
import base64
import numpy as np
from PIL import Image


class TACAnalyzer:
    @staticmethod
    def generate_previews(cmyk_bytes: bytes, width: int, height: int, max_preview_dim: int = 800):
        """
        Generates thumbnail previews:
        - Composite RGB image
        - Cyan isolation
        - Magenta isolation
        - Yellow isolation
        - Black isolation
        - TAC heatmap overlay
        Returns dictionary of base64 PNG data URLs.
        """
        # Downscale for fast visual rendering if large
        arr = np.frombuffer(cmyk_bytes, dtype=np.uint8).reshape((height, width, 4))
        scale = min(1.0, max_preview_dim / max(width, height))
        new_w = int(width * scale)
        new_h = int(height * scale)

        if scale < 1.0:
            pil_cmyk = Image.frombytes("CMYK", (width, height), cmyk_bytes)
            pil_cmyk = pil_cmyk.resize((new_w, new_h), Image.Resampling.BILINEAR)
            arr = np.array(pil_cmyk)

        # 1. Composite RGB
        # Standard CMYK to RGB formula
        c = arr[:, :, 0].astype(np.float32) / 255.0
        m = arr[:, :, 1].astype(np.float32) / 255.0
        y = arr[:, :, 2].astype(np.float32) / 255.0
        k = arr[:, :, 3].astype(np.float32) / 255.0

        r = (1.0 - c) * (1.0 - k) * 255.0
        g = (1.0 - m) * (1.0 - k) * 255.0
        b = (1.0 - y) * (1.0 - k) * 255.0
        rgb_arr = np.clip(np.dstack([r, g, b]), 0, 255).astype(np.uint8)
        img_composite = Image.fromarray(rgb_arr, "RGB")

        # 2. Individual Channel Separations (Visualized as colored inks)
        # Cyan: (1-c, 1-0.1*c, 1)
        c_sub = (arr[:, :, 0].astype(np.float32) * 0.1).astype(np.uint8)
        cyan_rgb = np.dstack([255 - arr[:, :, 0], 255 - c_sub, np.full((new_h, new_w), 255, dtype=np.uint8)])
        # Magenta: (1, 1-m, 1-0.2*m)
        m_sub = (arr[:, :, 1].astype(np.float32) * 0.2).astype(np.uint8)
        mag_rgb = np.dstack([np.full((new_h, new_w), 255, dtype=np.uint8), 255 - arr[:, :, 1], 255 - m_sub])
        # Yellow: (1, 1, 1-y)
        yel_rgb = np.dstack([np.full((new_h, new_w), 255, dtype=np.uint8), np.full((new_h, new_w), 255, dtype=np.uint8), 255 - arr[:, :, 2]])
        # Black: (1-k, 1-k, 1-k)
        blk_rgb = np.dstack([255 - arr[:, :, 3], 255 - arr[:, :, 3], 255 - arr[:, :, 3]])

        # 3. TAC Heatmap
        tac = (arr[:, :, 0].astype(np.int32) + arr[:, :, 1].astype(np.int32) +
               arr[:, :, 2].astype(np.int32) + arr[:, :, 3].astype(np.int32)) * 100 // 255

        # Colormap for TAC (0..400%)
        # 0-100: dark blue / transparent
        # 100-200: cyan -> green
        # 200-260: green -> yellow
        # 260-300: yellow -> orange (warning)
        # 300-340: orange -> red (high limit)
        # 340-400: deep red / magenta / purple (excessive!)
        heatmap_rgba = np.zeros((new_h, new_w, 4), dtype=np.uint8)

        # Vectorized color mapping
        mask_low = tac < 150
        heatmap_rgba[mask_low, 0] = 0
        heatmap_rgba[mask_low, 1] = (tac[mask_low] * 255 // 150).astype(np.uint8)
        heatmap_rgba[mask_low, 2] = 200
        heatmap_rgba[mask_low, 3] = (tac[mask_low] * 120 // 150).astype(np.uint8)

        mask_mid = (tac >= 150) & (tac < 260)
        t_mid = (tac[mask_mid] - 150).astype(np.float32) / 110.0
        heatmap_rgba[mask_mid, 0] = (t_mid * 255).astype(np.uint8)
        heatmap_rgba[mask_mid, 1] = 255
        heatmap_rgba[mask_mid, 2] = 0
        heatmap_rgba[mask_mid, 3] = 160

        mask_high = (tac >= 260) & (tac < 320)
        t_high = (tac[mask_high] - 260).astype(np.float32) / 60.0
        heatmap_rgba[mask_high, 0] = 255
        heatmap_rgba[mask_high, 1] = ((1.0 - t_high) * 165).astype(np.uint8)
        heatmap_rgba[mask_high, 2] = 0
        heatmap_rgba[mask_high, 3] = 200

        mask_crit = tac >= 320
        t_crit = np.clip((tac[mask_crit] - 320).astype(np.float32) / 80.0, 0, 1.0)
        heatmap_rgba[mask_crit, 0] = 255
        heatmap_rgba[mask_crit, 1] = 0
        heatmap_rgba[mask_crit, 2] = (t_crit * 255).astype(np.uint8)
        heatmap_rgba[mask_crit, 3] = 230

        img_cyan = Image.fromarray(cyan_rgb.astype(np.uint8), "RGB")
        img_mag = Image.fromarray(mag_rgb.astype(np.uint8), "RGB")
        img_yel = Image.fromarray(yel_rgb.astype(np.uint8), "RGB")
        img_blk = Image.fromarray(blk_rgb.astype(np.uint8), "RGB")
        img_tac = Image.fromarray(heatmap_rgba, "RGBA")

        def to_data_url(img, fmt="JPEG", quality=85):
            buf = io.BytesIO()
            if fmt == "JPEG":
                img.save(buf, format="JPEG", quality=quality)
                mime = "image/jpeg"
            else:
                img.save(buf, format="PNG")
                mime = "image/png"
            b64 = base64.b64encode(buf.getvalue()).decode("ascii")
            return f"data:{mime};base64,{b64}"

        return {
            "composite": to_data_url(img_composite, "JPEG", 90),
            "cyan": to_data_url(img_cyan, "JPEG", 85),
            "magenta": to_data_url(img_mag, "JPEG", 85),
            "yellow": to_data_url(img_yel, "JPEG", 85),
            "black": to_data_url(img_blk, "JPEG", 85),
            "tac_heatmap": to_data_url(img_tac, "PNG"),
            "preview_width": new_w,
            "preview_height": new_h
        }
