tensorview/colors.py

"""Colormap loading and management for CPT/ACT/RGB files."""

import os
import re
from typing import List, Tuple, Optional
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.colors import LinearSegmentedColormap, ListedColormap


def parse_cpt_file(filepath: str) -> LinearSegmentedColormap:
    """
    Load a GMT-style CPT (Color Palette Table) file.
    
    Args:
        filepath: Path to the CPT file
        
    Returns:
        matplotlib LinearSegmentedColormap
    """
    colors = []
    positions = []
    
    with open(filepath, 'r') as f:
        lines = f.readlines()
    
    # Parse CPT format
    for line in lines:
        line = line.strip()
        
        # Skip comments and empty lines
        if line.startswith('#') or not line or line.startswith('B') or line.startswith('F') or line.startswith('N'):
            continue
        
        parts = line.split()
        if len(parts) >= 8:
            # CPT format: z0 r0 g0 b0 z1 r1 g1 b1
            try:
                z0, r0, g0, b0, z1, r1, g1, b1 = map(float, parts[:8])
                
                # Normalize RGB values if they're in 0-255 range
                if r0 > 1 or g0 > 1 or b0 > 1:
                    r0, g0, b0 = r0/255, g0/255, b0/255
                if r1 > 1 or g1 > 1 or b1 > 1:
                    r1, g1, b1 = r1/255, g1/255, b1/255
                
                colors.extend([(r0, g0, b0), (r1, g1, b1)])
                positions.extend([z0, z1])
                
            except ValueError:
                continue
    
    if not colors:
        raise ValueError(f"No valid color data found in {filepath}")
    
    # Normalize positions to 0-1 range
    positions = np.array(positions)
    if len(set(positions)) > 1:
        positions = (positions - positions.min()) / (positions.max() - positions.min())
    else:
        positions = np.linspace(0, 1, len(positions))
    
    # Create colormap
    cmap_data = list(zip(positions, colors))
    cmap_data.sort(key=lambda x: x[0])  # Sort by position
    
    # Remove duplicates
    unique_data = []
    seen_positions = set()
    for pos, color in cmap_data:
        if pos not in seen_positions:
            unique_data.append((pos, color))
            seen_positions.add(pos)
    
    positions, colors = zip(*unique_data)
    
    # Create the colormap
    name = os.path.splitext(os.path.basename(filepath))[0]
    return LinearSegmentedColormap.from_list(name, list(zip(positions, colors)))


def parse_act_file(filepath: str) -> ListedColormap:
    """
    Load an Adobe Color Table (ACT) file.
    
    Args:
        filepath: Path to the ACT file
        
    Returns:
        matplotlib ListedColormap
    """
    with open(filepath, 'rb') as f:
        data = f.read()
    
    # ACT files are 768 bytes (256 colors * 3 RGB values)
    if len(data) != 768:
        raise ValueError(f"Invalid ACT file size: {len(data)} bytes (expected 768)")
    
    colors = []
    for i in range(0, 768, 3):
        r, g, b = data[i], data[i+1], data[i+2]
        colors.append((r/255.0, g/255.0, b/255.0))
    
    name = os.path.splitext(os.path.basename(filepath))[0]
    return ListedColormap(colors, name=name)


def parse_rgb_file(filepath: str) -> ListedColormap:
    """
    Load a simple RGB text file (one color per line).
    
    Args:
        filepath: Path to the RGB file
        
    Returns:
        matplotlib ListedColormap
    """
    colors = []
    
    with open(filepath, 'r') as f:
        lines = f.readlines()
    
    for line in lines:
        line = line.strip()
        
        # Skip comments and empty lines
        if line.startswith('#') or not line:
            continue
        
        # Parse RGB values
        parts = line.split()
        if len(parts) >= 3:
            try:
                r, g, b = map(float, parts[:3])
                
                # Normalize RGB values if they're in 0-255 range
                if r > 1 or g > 1 or b > 1:
                    r, g, b = r/255, g/255, b/255
                
                colors.append((r, g, b))
            except ValueError:
                continue
    
    if not colors:
        raise ValueError(f"No valid color data found in {filepath}")
    
    name = os.path.splitext(os.path.basename(filepath))[0]
    return ListedColormap(colors, name=name)


def load_cpt(filepath: str) -> LinearSegmentedColormap:
    """Load a CPT colormap file."""
    return parse_cpt_file(filepath)


def load_act(filepath: str) -> ListedColormap:
    """Load an ACT colormap file."""
    return parse_act_file(filepath)


def load_rgb(filepath: str) -> ListedColormap:
    """Load an RGB colormap file."""
    return parse_rgb_file(filepath)


def load_colormap(filepath: str) -> mcolors.Colormap:
    """
    Load a colormap from file, auto-detecting the format.
    
    Args:
        filepath: Path to the colormap file
        
    Returns:
        matplotlib Colormap
    """
    _, ext = os.path.splitext(filepath.lower())
    
    if ext == '.cpt':
        return load_cpt(filepath)
    elif ext == '.act':
        return load_act(filepath)
    elif ext in ['.rgb', '.txt']:
        return load_rgb(filepath)
    else:
        # Try to guess based on content
        try:
            return load_cpt(filepath)
        except:
            try:
                return load_rgb(filepath)
            except:
                raise ValueError(f"Cannot determine colormap format for {filepath}")


def get_matplotlib_colormaps() -> List[str]:
    """Get list of available matplotlib colormaps."""
    return sorted(plt.colormaps())


def create_diverging_colormap(name: str, colors: List[str], center: float = 0.5) -> LinearSegmentedColormap:
    """
    Create a diverging colormap.
    
    Args:
        name: Name for the colormap
        colors: List of color names/hex codes
        center: Position of the center color (0-1)
        
    Returns:
        LinearSegmentedColormap
    """
    return LinearSegmentedColormap.from_list(name, colors)


def reverse_colormap(cmap: mcolors.Colormap) -> mcolors.Colormap:
    """Reverse a colormap."""
    if hasattr(cmap, 'reversed'):
        return cmap.reversed()
    else:
        # For custom colormaps
        if isinstance(cmap, LinearSegmentedColormap):
            return LinearSegmentedColormap.from_list(
                f"{cmap.name}_r", 
                cmap(np.linspace(0, 1, 256))[::-1]
            )
        elif isinstance(cmap, ListedColormap):
            return ListedColormap(
                cmap.colors[::-1], 
                name=f"{cmap.name}_r"
            )
        else:
            return cmap


def colormap_to_cpt(cmap: mcolors.Colormap, filepath: str, n_colors: int = 256):
    """
    Save a matplotlib colormap as a CPT file.
    
    Args:
        cmap: matplotlib Colormap
        filepath: Output file path
        n_colors: Number of color levels
    """
    colors = cmap(np.linspace(0, 1, n_colors))
    
    with open(filepath, 'w') as f:
        f.write(f"# Color palette: {cmap.name}\n")
        f.write("# Created by TensorView\n")
        
        for i in range(len(colors) - 1):
            r0, g0, b0 = colors[i][:3]
            r1, g1, b1 = colors[i+1][:3]
            
            z0 = i / (n_colors - 1)
            z1 = (i + 1) / (n_colors - 1)
            
            f.write(f"{z0:.6f}\t{r0*255:.0f}\t{g0*255:.0f}\t{b0*255:.0f}\t")
            f.write(f"{z1:.6f}\t{r1*255:.0f}\t{g1*255:.0f}\t{b1*255:.0f}\n")


def get_colormap_info(cmap: mcolors.Colormap) -> dict:
    """Get information about a colormap."""
    info = {
        'name': getattr(cmap, 'name', 'unknown'),
        'type': type(cmap).__name__,
        'n_colors': getattr(cmap, 'N', 'continuous')
    }
    
    if isinstance(cmap, ListedColormap):
        info['n_colors'] = len(cmap.colors)
    
    return info