import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import io
import base64
from PIL import Image
from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg
import os
import yaml

# --- Config Loading ---
try:
    config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'config.yaml')
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
        viz_cfg = config.get('Visualize', {})
except Exception as e:
    print(f"[WARNING] Could not load config.yaml in plots.py: {e}")
    viz_cfg = {}

def render_heatmap(data, vmin=None, vmax=None, cmap='viridis'):
    """Converts a 2D numpy array to a colormapped PNG base64 string."""
    if data is None or data.size == 0:
        return None
    
    # Normalize to 0-1
    if vmin is None: vmin = np.min(data)
    if vmax is None: vmax = np.max(data)
    
    if vmax > vmin:
        norm_data = (data - vmin) / (vmax - vmin)
    else:
        norm_data = np.zeros_like(data)
    
    # Apply colormap
    color_mapped = matplotlib.colormaps[cmap](norm_data) # [H, W, 4]
    
    # Convert to 8-bit RGB
    rgb = (color_mapped[:, :, :3] * 255).astype(np.uint8)
    
    img = Image.fromarray(rgb)
    buffered = io.BytesIO()
    img.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("ascii")

class FastHeatmapEngine:
    """Stateful Matplotlib engine that reuses figure memory to achieve high-speed frame rendering."""
    def __init__(self, extent, cmap='jet', vmin=None, vmax=None, title='Heatmap', xlabel='X', ylabel='Y', xlim=None, ylim=None, aspect='auto', interpolation=None):
        self.vmin = vmin
        self.vmax = vmax
        
        # Load config-based overrides
        # We look for a key in Visualize that matches the plot type
        plot_type = 'rangeDoppler' if 'Doppler' in title else ('rangeAoA' if 'Azimuth' in title else None)
        if plot_type and plot_type in viz_cfg:
            cfg = viz_cfg[plot_type]
            if 'xRange' in cfg: xlim = cfg['xRange']
            if 'yRange' in cfg: ylim = cfg['yRange']

        self.fig = Figure(figsize=(6, 5), dpi=100)
        self.canvas = FigureCanvasAgg(self.fig)
        self.ax = self.fig.add_subplot(111)
        
        cm_obj = matplotlib.colormaps.get_cmap(cmap).copy()
        cm_obj.set_bad(color='white')
        
        # Dummy matrix to initialize geometry
        dummy = np.zeros((2, 2))
        self.im = self.ax.imshow(dummy, extent=extent, cmap=cm_obj, vmin=vmin, vmax=vmax, origin='upper', aspect=aspect, interpolation=interpolation)
        
        # Check for config overrides
        # (This is where the user's requested RD limit overrides would go)
        if xlim is not None: self.ax.set_xlim(xlim)
        if ylim is not None: self.ax.set_ylim(ylim)
        
        self.ax.set_xlabel(xlabel)
        self.ax.set_ylabel(ylabel)
        self.ax.set_title(title)
        self.ax.grid(color='gray', linestyle='--', linewidth=0.5, alpha=0.7)
        self.fig.colorbar(self.im, ax=self.ax, label='Magnitude (dB)')
        self.fig.tight_layout()
        
    def render(self, data):
        self.im.set_data(data)
        if self.vmin is None and self.vmax is None:
            v_low, v_high = np.nanmin(data), np.nanmax(data)
            self.im.set_clim(v_low if not np.isnan(v_low) else 0.0, v_high if not np.isnan(v_high) else 1.0)
            
        self.fig.canvas.draw()
        rgba = np.asarray(self.canvas.buffer_rgba())
        img = Image.fromarray(rgba)
        buf = io.BytesIO()
        img.save(buf, format='png')
        return base64.b64encode(buf.getvalue()).decode("ascii")

def postprocess_ra(ra_heatmap, range_axis, smooth_sigma=1.0):
    """
    Refined RA post-processing pipeline for Physical Realism.
    """
    # 1. Clutter removal (subtract per-range-bin mean to suppress static ground)
    clutter = np.mean(ra_heatmap, axis=1, keepdims=True)
    ra = ra_heatmap - (0.8 * clutter)
    ra = np.clip(ra, 1e-9, None)

    # 2. Physics-based dynamic range compression (Linear -> Log)
    SYSTEM_GAIN_OFFSET = 68.0 
    ra_db = 10 * np.log10(ra) - SYSTEM_GAIN_OFFSET
    
    # 3. Fixed dynamic range clipping (-5 to 45 dB)
    ra_db = np.clip(ra_db, -5, 45)

    # 4. Optional Gaussian smoothing
    if smooth_sigma > 0:
        from scipy.ndimage import gaussian_filter
        ra_db = gaussian_filter(ra_db, sigma=smooth_sigma)

    return ra_db

def scan_convert_ra(ra_heatmap, range_axis, angle_axis, img_size=512, max_display_range=None):
    """
    Polar-to-Cartesian scan conversion.
    Converts RA (Range, Angle) polar data into a Sector plot.
    """
    true_max_range = range_axis[-1]
    max_range = max_display_range if max_display_range is not None else true_max_range
    
    theta_min = angle_axis[0]
    theta_max = angle_axis[-1]

    # Create Cartesian Grid
    x = np.linspace(-max_range, max_range, img_size)
    y = np.linspace(max_range, 0, img_size) 
    X, Y = np.meshgrid(x, y)

    # Convert to Polar Coordinates
    R_query = np.sqrt(X**2 + Y**2)
    Theta_query = np.arctan2(X, Y)

    # Mask Valid Radar FOV
    fov_mask = (Theta_query >= theta_min) & (Theta_query <= theta_max) & (R_query <= max_range)

    # Map RA Heatmap to Cartesian Grid
    r_idx = np.clip(((R_query / true_max_range) * (ra_heatmap.shape[0] - 1)).astype(int), 0, ra_heatmap.shape[0] - 1)
    theta_range = theta_max - theta_min
    theta_idx = np.clip(((Theta_query - theta_min) / theta_range * (ra_heatmap.shape[1] - 1)).astype(int), 0, ra_heatmap.shape[1] - 1)

    # Project
    cartesian = np.full((img_size, img_size), np.nan, dtype=np.float64)
    cartesian[fov_mask] = ra_heatmap[r_idx[fov_mask], theta_idx[fov_mask]]

    return cartesian