#!/usr/bin/env python3
"""
V5: Velocity Magnitude Heatmap + Vertical Velocity Component
Shows WHERE flow is happening and whether it's UP or DOWN.
Better for coarse mesh data.
"""

import os
import re
import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import zoom
from pathlib import Path

try:
    import imageio
    HAS_IMAGEIO = True
except ImportError:
    HAS_IMAGEIO = False

plt.rcParams['font.family'] = 'sans-serif'


def parse_openfoam_field(filepath, field_type='scalar'):
    with open(filepath, 'r') as f:
        content = f.read()
    if 'internalField' not in content:
        return None
    section = content.split('internalField', 1)[1]
    if 'boundaryField' in section:
        section = section.split('boundaryField', 1)[0]
    section = section.strip()
    if section.startswith('uniform'):
        match = re.search(r'uniform\s+([^;]+);', section)
        if match:
            val_str = match.group(1).strip()
            if field_type == 'scalar':
                return float(val_str)
    elif section.startswith('nonuniform'):
        start_idx = section.find('(')
        if start_idx == -1:
            return None
        list_content = section[start_idx:]
        if field_type == 'scalar':
            floats = re.findall(r'[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?', list_content)
            return np.array([float(f) for f in floats])
        elif field_type == 'vector':
            pattern = r'\(\s*([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)\s+([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)\s+([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)\s*\)'
            matches = re.findall(pattern, section)
            return np.array([(float(x), float(y), float(z)) for x, y, z in matches])
    return None


def get_timesteps(case_path):
    timesteps = []
    for item in os.listdir(case_path):
        if os.path.isdir(os.path.join(case_path, item)):
            try:
                timesteps.append((float(item), item))
            except ValueError:
                pass
    return sorted(timesteps, key=lambda x: x[0])


def create_v5_visualization(case_path, output_path, fps=8):
    """V5: Heatmap of vertical velocity (Uy) showing up vs down flow."""
    
    print(f"V5 Visualization: {case_path}")
    
    nx, ny = 20, 40
    zoom_factor = 10
    
    x_high = np.linspace(0, 10, nx * zoom_factor)
    y_high = np.linspace(0, 20, ny * zoom_factor)
    X_high, Y_high = np.meshgrid(x_high, y_high)
    
    timesteps = get_timesteps(case_path)
    if not timesteps:
        print("No timesteps!")
        return
    
    bg_color = '#0d1117'
    plt.style.use('dark_background')
    
    frames = []
    
    for i, (t, t_dir) in enumerate(timesteps):
        fig, ax = plt.subplots(figsize=(10, 12), dpi=120)
        fig.patch.set_facecolor(bg_color)
        ax.set_facecolor(bg_color)
        
        alpha_path = os.path.join(case_path, t_dir, 'alpha.water')
        U_path = os.path.join(case_path, t_dir, 'U')
        
        if not os.path.exists(alpha_path):
            plt.close()
            continue
        
        alpha = parse_openfoam_field(alpha_path, 'scalar')
        if isinstance(alpha, float):
            alpha = np.full(nx * ny, alpha)
        alpha_2d = alpha.reshape((ny, nx))
        
        if os.path.exists(U_path):
            U = parse_openfoam_field(U_path, 'vector')
        else:
            U = None
        
        if U is not None and len(U) == nx * ny:
            Uy = np.array([v[1] for v in U]).reshape((ny, nx))
        else:
            Uy = np.zeros((ny, nx))
        
        # Interpolate
        alpha_high = zoom(alpha_2d, zoom_factor, order=3)
        Uy_high = zoom(Uy, zoom_factor, order=3)
        
        # Mask gas phase
        Uy_liquid = np.where(alpha_high > 0.3, Uy_high, np.nan)
        
        # === VISUALIZATION ===
        
        # 1. Vertical velocity heatmap (blue = down, red = up)
        # Scale in mm/s for readability
        Uy_mm = Uy_liquid * 1000  # m/s -> mm/s
        
        vmax = 0.02  # mm/s symmetric scale
        im = ax.imshow(Uy_mm, extent=[0, 10, 0, 20], origin='lower',
                      cmap='RdBu_r', vmin=-vmax, vmax=vmax, aspect='equal')
        
        # 2. Interface contour
        ax.contour(X_high, Y_high, alpha_high, levels=[0.5],
                  colors=['#00ff88'], linewidths=2)
        
        # 3. Colorbar
        cbar = plt.colorbar(im, ax=ax, shrink=0.6, pad=0.02)
        cbar.set_label('Vertical Velocity (mm/s)\n⬆ UP (red) | ⬇ DOWN (blue)', 
                      color='white', fontsize=10)
        cbar.ax.tick_params(colors='white')
        
        # 4. Annotations
        ax.plot([4, 6], [19.7, 19.7], color='#ff6b6b', linewidth=6, solid_capstyle='round')
        ax.text(5, 18.5, 'HOT ZONE', color='#ff6b6b', fontsize=11, ha='center', fontweight='bold')
        
        # Gravity
        ax.annotate('', xy=(9.3, 1), xytext=(9.3, 3.5),
                   arrowprops=dict(arrowstyle='->', color='white', lw=2))
        ax.text(9.3, 0.5, 'g', color='white', fontsize=12, ha='center', fontweight='bold')
        
        # Key insight
        if t > 0.02:
            ax.text(5, 8, 'RED = Liquid flowing UP\n(against gravity!)',
                   color='white', fontsize=11, ha='center',
                   bbox=dict(facecolor='#1a1a2e', edgecolor='#00ff88', 
                            alpha=0.9, boxstyle='round,pad=0.5'))
        
        # Axes
        ax.set_xlim(0, 10)
        ax.set_ylim(0, 20)
        ax.set_xlabel('Width (mm)', color='white', fontsize=11)
        ax.set_ylabel('Height (mm)', color='white', fontsize=11)
        ax.tick_params(colors='white')
        
        ax.set_title(f'Marangoni Flow — Vertical Velocity Field\nt = {t*1000:.0f} ms',
                    color='white', fontsize=14, fontweight='bold', pad=10)
        
        for spine in ax.spines.values():
            spine.set_color('#30363d')
        
        plt.tight_layout()
        
        frame_path = os.path.join(output_path, f'frame_v5_{i:04d}.png')
        plt.savefig(frame_path, dpi=120, facecolor=bg_color, bbox_inches='tight')
        frames.append(frame_path)
        plt.close()
        
        print(f"  Frame {i+1}/{len(timesteps)}: t = {t*1000:.0f} ms")
    
    if HAS_IMAGEIO and frames:
        print("Creating V5 GIF...")
        images = [imageio.imread(f) for f in frames]
        for _ in range(4):
            images.append(images[-1])
        gif_path = os.path.join(output_path, 'marangoni_v5.gif')
        imageio.mimsave(gif_path, images, fps=fps, loop=0)
        print(f"✅ V5 saved: {gif_path}")
    
    return frames


if __name__ == '__main__':
    base_dir = Path(__file__).parent.parent
    case_path = base_dir / 'cfd_output'
    output_path = base_dir / 'marangoni_video_v5'
    
    if output_path.exists():
        import shutil
        shutil.rmtree(output_path)
    output_path.mkdir()
    
    create_v5_visualization(str(case_path), str(output_path))