
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import yaml
import math
from pathlib import Path
from ase.io import read

def plot_md_log(csv_path: str, out_prefix: str = "md_plot"):
    """Plot T, P, E from md.csv."""
    if not Path(csv_path).exists():
        print(f"Error: {csv_path} not found.")
        return
    
    df = pd.read_csv(csv_path)
    fig, axes = plt.subplots(3, 1, figsize=(10, 12), sharex=True)
    
    # Temperature
    axes[0].plot(df['time_fs'], df['T_K'], color='r')
    axes[0].set_ylabel("Temperature (K)")
    
    # Potential Energy
    axes[1].plot(df['time_fs'], df['Epot_eV'], color='b')
    axes[1].set_ylabel("Potential Energy (eV)")
    
    # Pressure
    if 'P_GPa' in df.columns:
        axes[2].plot(df['time_fs'], df['P_GPa'], color='g')
        axes[2].set_ylabel("Pressure (GPa)")
    
    axes[2].set_xlabel("Time (fs)")
    plt.tight_layout()
    plt.savefig(f"{out_prefix}.pdf")
    print(f"MD plots saved to {out_prefix}.pdf")

def plot_rdf(traj_path: str, out_prefix: str = "rdf_plot", r_max: float = 10.0, n_bins: int = 200):
    """Calculate and plot Radial Distribution Function (RDF) from trajectory."""
    print(f"Reading trajectory for RDF: {traj_path}")
    try:
        # Read the last 20% of frames for better statistics of equilibrated state
        all_configs = read(traj_path, index=':')
        n_frames = len(all_configs)
        start_idx = int(n_frames * 0.8)
        configs = all_configs[start_idx:]
        print(f"Using last {len(configs)} frames for RDF averaging.")
    except Exception as e:
        print(f"Error reading trajectory: {e}")
        return

    from ase.geometry.analysis import Analysis
    
    # Get all unique elements
    species = sorted(list(set(configs[0].get_chemical_symbols())))
    pairs = []
    for i, s1 in enumerate(species):
        for j, s2 in enumerate(species):
            if i <= j:
                pairs.append((s1, s2))

    plt.figure(figsize=(10, 6))
    
    for s1, s2 in pairs:
        print(f"  - Calculating RDF for {s1}-{s2}...")
        all_rdf = []
        for atoms in configs:
            ana = Analysis(atoms)
            # rdf returns [r_values, rdf_values]
            rdf = ana.get_rdf(rmax=r_max, nbins=n_bins, elements=[s1, s2])[0]
            all_rdf.append(rdf)
        
        avg_rdf = np.mean(all_rdf, axis=0)
        r_axis = np.linspace(0, r_max, n_bins)
        plt.plot(r_axis, avg_rdf, label=f"{s1}-{s2}")

    plt.xlabel(r"Distance $r$ ($\AA$)")
    plt.ylabel(r"Radial Distribution $g(r)$")
    plt.title(f"Radial Distribution Function (Averaged over last {len(configs)} frames)")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.xlim(0, r_max)
    plt.ylim(bottom=0)
    plt.tight_layout()
    plt.savefig(f"{out_prefix}.pdf")
    print(f"RDF plot saved to {out_prefix}.pdf")

def _parse_phonon_dat(dat_path):
    """
    Parse phonopy-style .dat files. Returns a list of bands, 
    where each band is a list of segments, and each segment is a list of [x, y, (g)].
    """
    if not Path(dat_path).exists():
        print(f"Error: File not found: {dat_path}")
        return []

    bands = []
    current_band = []
    current_segment = []
    
    with open(dat_path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                if current_segment:
                    current_band.append(np.array(current_segment))
                    current_segment = []
                continue
            if line.startswith('#'):
                if "mode" in line.lower() and current_band:
                    bands.append(current_band)
                    current_band = []
                continue
            
            parts = [float(x) for x in line.split()]
            current_segment.append(parts)
            
    if current_segment:
        current_band.append(np.array(current_segment))
    if current_band:
        bands.append(current_band)
        
    return bands

def _parse_phonon_yaml(yaml_path):
    """
    Parse phonopy band.yaml or gruneisen.yaml.
    Returns: (bands, labels_dict)
    labels_dict: {'ticks': [dist1, dist2, ...], 'labels': ['G', 'M', ...]}
    """
    if not Path(yaml_path).exists():
        return [], None

    print(f"Parsing YAML for data and labels: {yaml_path}")
    try:
        with open(yaml_path, 'r') as f:
            try:
                from yaml import CLoader as Loader
            except ImportError:
                from yaml import Loader
            data = yaml.load(f, Loader=Loader)
    except Exception as e:
        print(f"Error reading YAML: {e}")
        return [], None

    if 'phonon' not in data:
        return [], None

    nq = len(data['phonon'])
    nband = len(data['phonon'][0]['band'])
    distances = np.array([q['distance'] for q in data['phonon']])
    
    segment_indices = [0]
    if 'segment_nqpoint' in data:
        curr = 0
        for n in data['segment_nqpoint']:
            curr += n
            segment_indices.append(curr)
    else:
        for i in range(1, nq):
            if distances[i] < distances[i-1]:
                segment_indices.append(i)
        segment_indices.append(nq)

    bands = []
    for b_idx in range(nband):
        curr_band = []
        for s_idx in range(len(segment_indices)-1):
            start, end = segment_indices[s_idx], segment_indices[s_idx+1]
            seg_data = []
            for i in range(start, end):
                q = data['phonon'][i]
                row = [q['distance'], q['band'][b_idx]['frequency']]
                if 'gruneisen' in q['band'][b_idx]:
                    row.append(q['band'][b_idx]['gruneisen'])
                seg_data.append(row)
            curr_band.append(np.array(seg_data))
        bands.append(curr_band)

    labels_info = None
    if 'labels' in data and 'segment_nqpoint' in data:
        ticks = []
        labels = []
        curr_idx = 0
        for i, pair in enumerate(data['labels']):
            if i == 0 or data['labels'][i-1][1] != pair[0]:
                ticks.append(distances[curr_idx])
                labels.append(pair[0])
            curr_idx += data['segment_nqpoint'][i]
            actual_idx = min(curr_idx - 1, nq - 1)
            ticks.append(distances[actual_idx])
            labels.append(pair[1])
            
        unique_ticks = []
        unique_labels = []
        for t, l in zip(ticks, labels):
            if not unique_ticks or not np.isclose(t, unique_ticks[-1]):
                unique_ticks.append(t)
                unique_labels.append(l)
            else:
                if l != unique_labels[-1]:
                    unique_labels[-1] = f"{unique_labels[-1]}|{l}"
        labels_info = {'ticks': unique_ticks, 'labels': unique_labels}

    return bands, labels_info

def plot_phonon_band(dat_path, out_pdf="phonon_band.pdf", fmin=None, fmax=None, labels=None, yaml_path=None):
    """Plot phonon dispersion from .dat or .yaml file."""
    labels_info = None
    
    # 1. Load Data
    if str(dat_path).endswith('.yaml'):
        bands, labels_info = _parse_phonon_yaml(dat_path)
    else:
        bands = _parse_phonon_dat(dat_path)
    
    # 2. Load Labels from YAML (if not already loaded from dat_path)
    if labels_info is None:
        y_path = Path(yaml_path) if yaml_path else Path(dat_path).with_suffix('.yaml')
        if y_path.exists():
            _, labels_info = _parse_phonon_yaml(y_path)

    if not bands:
        print(f"Error: No data found in {dat_path}")
        return

    plt.figure(figsize=(8, 6))
    plt.rcParams["pdf.fonttype"] = 42
    plt.rcParams["font.family"] = "serif"
    
    all_x = []
    segment_boundaries = [0]
    
    for i, band in enumerate(bands):
        curr_x = 0
        for j, segment in enumerate(band):
            x = segment[:, 0]
            y = segment[:, 1]
            plt.plot(x, y, color='b', linewidth=1.0, alpha=0.8)
            if i == 0:
                all_x.extend(x)
                curr_x = x[-1]
                segment_boundaries.append(curr_x)

    all_x = np.unique(all_x)
    plt.xlim(all_x[0], all_x[-1])
    if fmin is not None or fmax is not None:
        plt.ylim(fmin, fmax)
    
    plt.ylabel("Frequency (THz)")
    plt.axhline(0, color='k', linestyle='-', linewidth=0.5)
    
    unique_boundaries = np.unique(segment_boundaries)
    if labels:
        label_list = labels.split()
        if len(label_list) == len(unique_boundaries):
            clean_labels = [l.replace("GAMMA", "$\\Gamma$").replace("GM", "$\\Gamma$") for l in label_list]
            plt.xticks(unique_boundaries, clean_labels)
        else:
            print(f"Warning: Number of labels ({len(label_list)}) does not match boundaries ({len(unique_boundaries)})")
            plt.xticks(unique_boundaries)
    elif labels_info:
        clean_labels = [l.replace("GAMMA", "$\\Gamma$").replace("GM", "$\\Gamma$") for l in labels_info['labels']]
        plt.xticks(labels_info['ticks'], clean_labels)
        unique_boundaries = labels_info['ticks']
    else:
        plt.xticks(unique_boundaries)

    for b in unique_boundaries:
        plt.axvline(b, color='k', linestyle='-', linewidth=0.5, alpha=0.5)

    plt.tight_layout()
    plt.savefig(out_pdf)
    plt.close()
    print(f"Phonon band plot saved to {out_pdf}")

def plot_gruneisen_band(dat_path, out_prefix="gruneisen", fmin=None, fmax=None, gmin=None, gmax=None, filter_outliers=3.0, labels=None, yaml_path=None):
    """Plot Gruneisen band and mode parameters from .dat or .yaml file."""
    labels_info = None
    
    # 1. Load Data
    if str(dat_path).endswith('.yaml'):
        bands, labels_info = _parse_phonon_yaml(dat_path)
    else:
        bands = _parse_phonon_dat(dat_path)
    
    # 2. Load Labels from YAML
    if labels_info is None:
        y_path = Path(yaml_path) if yaml_path else Path(dat_path).with_suffix('.yaml')
        if y_path.exists():
            _, labels_info = _parse_phonon_yaml(y_path)

    if not bands:
        print(f"Error: No data found in {dat_path}")
        return

    all_data = []
    for band in bands:
        for segment in band:
            all_data.append(segment)
    data = np.vstack(all_data)
    
    x = data[:, 0]
    y = data[:, 1]
    if data.shape[1] < 3:
        print(f"Error: No Grüneisen data found in {dat_path}. Make sure the file contains 3 columns.")
        return
    g = data[:, 2]
    
    # Determine plot limits first
    # If user provided limits, use them. Otherwise, calculate using outlier-aware defaults.
    q1, q3 = np.percentile(g, [25, 75])
    iqr = q3 - q1
    auto_vmin = q1 - filter_outliers * iqr
    auto_vmax = q3 + filter_outliers * iqr
    
    # Check for significant negative values (tolerance -0.05)
    has_significant_negative = g.min() < -0.05
    
    if gmin is None:
        if not has_significant_negative:
            # Case: Effectively positive -> Start from 0
            vmin = 0.0
        else:
            # Case: Significant negative -> Use auto limit
            vmin = float(math.floor(max(g.min(), auto_vmin)))
    else:
        vmin = gmin

    vmax = gmax if gmax is not None else float(math.ceil(min(g.max(), auto_vmax)))
    
    # Data mask: Only remove points that are truly outside our determined plot range
    mask = (g >= vmin) & (g <= vmax)
    print(f"Plotting range: [{vmin:.2f}, {vmax:.2f}]. Points outside this range are hidden.")
    
    x_filt, y_filt, g_filt = x[mask], y[mask], g[mask]
    
    # 1. Gruneisen Band Plot
    fig, ax = plt.subplots(figsize=(8, 6))
    plt.rcParams["pdf.fonttype"] = 42
    
    # Use TwoSlopeNorm to keep 0 as white (diverging scale)
    from matplotlib.colors import TwoSlopeNorm, Normalize, LinearSegmentedColormap
    
    # Force range to include 0 to ensure white center
    vmin = min(vmin, 0.0)
    vmax = max(vmax, 0.0)
    
    # Default cmap
    cmap_to_use = 'bwr'
    
    if vmin == 0 and vmax == 0:
        # Trivial case: all 0
        norm = Normalize(vmin=-1, vmax=1)
        ticks = [0]
    elif vmin == 0:
        # Positive only: Truncate bwr to White->Red (0.5 to 1.0)
        original_cmap = plt.get_cmap('bwr')
        colors = original_cmap(np.linspace(0.5, 1.0, 128))
        cmap_to_use = LinearSegmentedColormap.from_list('bwr_upper', colors)
        
        norm = Normalize(vmin=0, vmax=vmax)
        n_side = 4
        ticks = np.linspace(0, vmax, n_side + 1)
    elif vmax == 0:
        # Negative only: Truncate bwr to Blue->White (0.0 to 0.5)
        original_cmap = plt.get_cmap('bwr')
        colors = original_cmap(np.linspace(0.0, 0.5, 128))
        cmap_to_use = LinearSegmentedColormap.from_list('bwr_lower', colors)
        
        norm = Normalize(vmin=vmin, vmax=0)
        n_side = 4
        ticks = np.linspace(vmin, 0, n_side + 1)
    else:
        # Crossing 0: TwoSlopeNorm for full dynamic range on both sides
        norm = TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax)
        n_side = 4
        
        ticks_neg = np.linspace(vmin, 0, n_side + 1) if vmin < 0 else []
        ticks_pos = np.linspace(0, vmax, n_side + 1) if vmax > 0 else []
        
        # Combine and ensure unique 0
        if len(ticks_neg) > 0 and len(ticks_pos) > 0:
             ticks = np.concatenate([ticks_neg, ticks_pos[1:]])
        elif len(ticks_neg) > 0:
             ticks = ticks_neg
        else:
             ticks = ticks_pos

    sc = ax.scatter(x_filt, y_filt, c=g_filt, cmap=cmap_to_use, s=5, norm=norm, alpha=0.8)
    cb = fig.colorbar(sc, ax=ax, ticks=ticks)
    cb.set_label("Grüneisen parameter")
    
    ax.set_ylabel("Frequency (THz)")
    ax.set_xlim(x.min(), x.max())
    if fmin is not None or fmax is not None: ax.set_ylim(fmin, fmax)
    
    segment_boundaries = [0]
    if len(bands) > 0 and len(bands[0]) > 0:
        for seg in bands[0]:
            segment_boundaries.append(seg[-1, 0])
    unique_boundaries = np.unique(segment_boundaries)
    
    if labels:
        label_list = labels.split()
        if len(label_list) == len(unique_boundaries):
            clean_labels = [l.replace("GAMMA", "$\\Gamma$").replace("GM", "$\\Gamma$") for l in label_list]
            ax.set_xticks(unique_boundaries)
            ax.set_xticklabels(clean_labels)
    elif labels_info:
        clean_labels = [l.replace("GAMMA", "$\\Gamma$").replace("GM", "$\\Gamma$") for l in labels_info['labels']]
        ax.set_xticks(labels_info['ticks'])
        ax.set_xticklabels(clean_labels)
        unique_boundaries = labels_info['ticks']

    for b in unique_boundaries:
        ax.axvline(b, color='k', linestyle='-', linewidth=0.5, alpha=0.5)

    plt.tight_layout()
    plt.savefig(f"{out_prefix}_band.pdf")
    print(f"Gruneisen band plot saved to {out_prefix}_band.pdf")
    
    # 2. Mode Parameter Plot
    fig, ax = plt.subplots(figsize=(8, 6))
    
    # Apply Piecewise Linear Scale to Y-axis if 0 is within range (and range is non-trivial)
    if vmin < 0 < vmax:
        from matplotlib.scale import FuncScale
        def forward(y):
            return np.where(y > 0, y / vmax, y / abs(vmin))
        def inverse(y):
            return np.where(y > 0, y * vmax, y * abs(vmin))
        ax.set_yscale('function', functions=(forward, inverse))
        ax.set_yticks(ticks)
    
    ax.scatter(x_filt, g_filt, c=g_filt, cmap=cmap_to_use, s=5, norm=norm, alpha=0.8)
    cb2 = fig.colorbar(sc, ax=ax, ticks=ticks)
    cb2.set_label("Grüneisen parameter")
    
    ax.set_ylabel("Mode Grüneisen parameter")
    ax.set_xlim(x.min(), x.max())
    ax.set_ylim(vmin, vmax)
    ax.axhline(0, color='gray', linestyle=':', alpha=0.5)
    
    for b in unique_boundaries:
        ax.axvline(b, color='k', linestyle='-', linewidth=0.5, alpha=0.5)
    
    if labels or labels_info:
        ax.set_xticks(unique_boundaries)
        ax.set_xticklabels(clean_labels)

    plt.tight_layout()
    plt.savefig(f"{out_prefix}_mode_parameter.pdf")
    print(f"Mode Gruneisen plot saved to {out_prefix}_mode_parameter.pdf")
    plt.close('all')

def _get_bz_wireframe(cell):
    """
    Generate Brillouin Zone wireframe for a given real-space cell.
    Returns: list of lines (pairs of 3D points).
    """
    from scipy.spatial import Voronoi, ConvexHull
    
    # Reciprocal cell (b vectors as rows, without 2pi factor from ase)
    # We multiply by 2pi to match physics convention
    rec_cell = cell.reciprocal() * 2 * np.pi
    
    # Generate neighbor lattice points (3x3x3 grid is usually enough for 1st BZ)
    nx, ny, nz = np.meshgrid([-1, 0, 1], [-1, 0, 1], [-1, 0, 1])
    points = np.dot(np.stack([nx.flatten(), ny.flatten(), nz.flatten()], axis=1), rec_cell)
    
    vor = Voronoi(points)
    
    # Center is at 0,0,0 (index 13 in the flattened 3x3x3 grid)
    center_idx = 13 
    region_idx = vor.point_region[center_idx]
    region = vor.regions[region_idx]
    
    # Vertices of the BZ
    vertices = vor.vertices[region]
    
    # Use ConvexHull to get edges of the polyhedron
    hull = ConvexHull(vertices)
    
    wireframe = []
    # Collect unique edges to avoid drawing duplicates
    edges = set()
    for simplex in hull.simplices:
        # simplex is a face (list of vertex indices)
        # We add edges (v0-v1, v1-v2, ...)
        for i in range(len(simplex)):
            v1, v2 = sorted((simplex[i], simplex[(i+1)%len(simplex)]))
            edges.add((v1, v2))
            
    for v1, v2 in edges:
        wireframe.append((vertices[v1], vertices[v2]))
            
    return wireframe

def plot_gruneisen_3d(dat_path, poscar_path=None, out_html="gruneisen_3d.html", gmin=None, gmax=None, filter_outliers=3.0, mode="mean"):
    """
    Generates a 3D Plotly visualization of Gruneisen parameters on the Brillouin Zone.
    """
    try:
        import plotly.graph_objects as go
        from scipy.spatial import ConvexHull, KDTree
    except ImportError:
        print("Error: 'plotly' and 'scipy' are required for 3D visualization.")
        return

    # Metadata placeholders
    meta_formula = "N/A"
    meta_sg = "N/A"
    meta_path = "N/A"

    # 1. Load Data & Parse Header
    try:
        # Read header metadata
        with open(dat_path, 'r') as f:
            for line in f:
                if line.startswith("#"):
                    if "Formula:" in line: meta_formula = line.split(":", 1)[1].strip()
                    if "SpaceGroup:" in line: meta_sg = line.split(":", 1)[1].strip()
                    if "Path:" in line: meta_path = line.split(":", 1)[1].strip()
                else:
                    break # Stop at data

        # Columns: qx qy qz weight band_index frequency gruneisen_param
        # Using pandas for easier grouping
        import pandas as pd
        df = pd.read_csv(dat_path, sep=r'\s+', comment='#', names=['qx', 'qy', 'qz', 'weight', 'band_idx', 'freq', 'gamma'])
        
        # Drop NaNs
        df = df.dropna()
        
        if df.empty:
            print(f"Error: No valid data found in {dat_path} (after removing NaNs).")
            return
            
    except Exception as e:
        print(f"Error loading {dat_path}: {e}")
        return

    # Aggregate data: Mean and Min per q-point
    q_groups = df.groupby(['qx', 'qy', 'qz'])['gamma'].agg(['mean', 'min']).reset_index()
    
    q_frac_base = q_groups[['qx', 'qy', 'qz']].values
    
    # Legacy aliases for intermediate blocks (Geometry/Symmetry)
    q_frac = q_frac_base
    if mode == "min":
        points_g = q_groups['min'].values
    else:
        points_g = q_groups['mean'].values
    
    # 2. Geometry & BZ
    if poscar_path is None:
        # Try to infer structure file: POSCAR -> POSCAR-input
        parent_dir = Path(dat_path).parent
        candidates = ["POSCAR", "POSCAR-input"]
        
        for c in candidates:
            cand_path = parent_dir / c
            if cand_path.exists():
                poscar_path = cand_path
                break
        
        if poscar_path is None:
            print("Warning: Neither 'POSCAR' nor 'POSCAR-input' found. Cannot generate BZ wireframe. Plotting in fractional coordinates.")

    wireframe = []
    special_points_cart = {}
    special_points_frac = {}
    special_points_frac_vals = []
    
    # Initialize metadata (will be overwritten if POSCAR is loaded or by file header)
    if meta_formula == "N/A": meta_formula = "N/A"
    if meta_sg == "N/A": meta_sg = "N/A" 
    if meta_path == "N/A": meta_path = "N/A"

    if poscar_path:
        try:
            from ase.io import read
            structure = read(str(poscar_path))
            # Update formula if still N/A
            if meta_formula == "N/A": meta_formula = structure.get_chemical_formula()
            
            cell_tuple = (structure.get_cell(), structure.get_scaled_positions(), structure.get_atomic_numbers())
            
            # --- BZ Wireframe ---
            from macer.utils.viz_tools import _get_bz_wireframe
            wireframe = _get_bz_wireframe(structure.get_cell())
            rec_cell = structure.get_cell().reciprocal() * 2 * np.pi

            # --- Special Points & Path ---
            try:
                try:
                    import seekpath
                except ImportError:
                    import macer.externals.seekpath_bundled as seekpath
                
                res = seekpath.get_explicit_k_path(cell_tuple)
                
                # Extract Path
                if meta_path == "N/A" and 'path' in res:
                    paths = res['path']
                    flat_path = []
                    for p in paths:
                        start, end = p
                        if not flat_path:
                            flat_path.append(start)
                            flat_path.append(end)
                        else:
                            if flat_path[-1] == start:
                                flat_path.append(end)
                            else:
                                flat_path.append(f"| {start}")
                                flat_path.append(end)
                    meta_path = "-".join([l.replace("GAMMA", "Γ").replace("GM", "Γ") for l in flat_path])

                # res['point_coords'] is {Label: [x, y, z]} (fractional)
                special_points_frac_vals = list(res['point_coords'].values())
                for label, coords in res['point_coords'].items():
                    cart_coord = np.dot(coords, rec_cell)
                    clean_label = label.replace("GAMMA", "Γ").replace("GM", "Γ")
                    special_points_cart[clean_label] = cart_coord
            except Exception as e:
                print(f"  Warning: Could not determine special points: {e}")

            # --- Symmetry Expansion ---
            try:
                try:
                    import spglib
                except ImportError:
                    import macer.externals.spglib_bundled as spglib
                
                dataset = spglib.get_symmetry(cell_tuple)
                rotations = dataset['rotations'] # (N_op, 3, 3)
                
                # Update Space Group if missing
                if meta_sg == "N/A":
                    try:
                        meta_sg = spglib.get_spacegroup(cell_tuple, symprec=1e-5)
                    except: pass
                
                # Expand q-points
                # q_new = R @ q (Generate star of k)
                expanded_q = []
                expanded_g = []
                
                for q, g_val in zip(q_frac, points_g):
                    # rotations: (N, 3, 3), q: (3,) -> Result: (N, 3)
                    qs_star = np.matmul(rotations, q)
                    
                    # Unique-ify equivalent points in the star
                    # Rounding to handle noise
                    qs_star_unique = np.unique(np.round(qs_star, decimals=8), axis=0)
                    
                    expanded_q.append(qs_star_unique)
                    expanded_g.append(np.full(len(qs_star_unique), g_val))
                
                if expanded_q:
                    q_frac = np.vstack(expanded_q)
                    points_g = np.concatenate(expanded_g)
                    
                    # --- Map to 1st BZ (Wigner-Seitz cell) ---
                    # Find integer shift n such that |q_cart + G_n| is minimized (closest to Gamma)
                    # Shifts grid: -1, 0, 1 (sufficient for 1st BZ)
                    shifts = np.array(np.meshgrid([-1, 0, 1], [-1, 0, 1], [-1, 0, 1])).T.reshape(-1, 3)
                    shift_cart = np.dot(shifts, rec_cell) # (27, 3)
                    
                    q_cart_base = np.dot(q_frac, rec_cell) # (N, 3)
                    
                    # (N, 1, 3) + (1, 27, 3) -> (N, 27, 3)
                    candidates = q_cart_base[:, np.newaxis, :] + shift_cart[np.newaxis, :, :]
                    dists = np.linalg.norm(candidates, axis=2) # (N, 27)
                    best_indices = np.argmin(dists, axis=1)
                    
                    best_shifts = shifts[best_indices]
                    q_frac = q_frac + best_shifts
                    
                    # --- Convex Hull Filtering (DEFAULT) ---
                    print("  Filtering q-points by Convex Hull of Special Points...")
                    
                    if special_points_cart:
                        try:
                            from scipy.spatial import Delaunay
                            # Create hull from special points (cartesian)
                            # Ensure Gamma is included for robustness, though usually it is.
                            hull_pts = np.array(list(special_points_cart.values()))
                            
                            # Check if hull is 3D
                            if len(hull_pts) >= 4:
                                hull = Delaunay(hull_pts)
                                
                                # We need q in Cartesian for the check
                                q_cart_check = np.dot(q_frac, rec_cell)
                                
                                # find_simplex returns -1 if point is outside
                                inside_mask = hull.find_simplex(q_cart_check) >= 0
                                
                                prev_len_h = len(q_frac)
                                q_frac = q_frac[inside_mask]
                                points_g = points_g[inside_mask]
                                
                                print(f"  Convex Hull Filter: {prev_len_h} -> {len(q_frac)} points inside IBZ-like region.")
                            else:
                                print("  Warning: Not enough special points to form a 3D hull. Skipping filtering.")
                        except Exception as e:
                            print(f"  Error during Convex Hull filtering: {e}")
                    else:
                        print("  Warning: No special points found. Cannot filter by Convex Hull.")

            except ImportError:
                print("  Warning: spglib not found. Skipping symmetry expansion.")
            except Exception as e:
                print(f"  Warning: Symmetry expansion failed: {e}")

            # --- Special Points (Re-calc for safety or use existing) ---
            try:
                try:
                    import seekpath
                except ImportError:
                    import macer.externals.seekpath_bundled as seekpath
                
                res = seekpath.get_explicit_k_path(cell_tuple)
                
                # Update Path if missing
                if meta_path == "N/A" and 'path' in res:
                    paths = res['path']
                    flat_path = []
                    for p in paths:
                        start, end = p
                        if not flat_path:
                            flat_path.append(start)
                            flat_path.append(end)
                        else:
                            if flat_path[-1] == start:
                                flat_path.append(end)
                            else:
                                flat_path.append(f"| {start}")
                                flat_path.append(end)
                    meta_path = "-".join([l.replace("GAMMA", "Γ").replace("GM", "Γ") for l in flat_path])

                # res['point_coords'] is {Label: [x, y, z]} (fractional)
                special_points_frac = {}
                for label, coords in res['point_coords'].items():
                    cart_coord = np.dot(coords, rec_cell)
                    clean_label = label.replace("GAMMA", "Γ").replace("GM", "Γ")
                    special_points_cart[clean_label] = cart_coord
                    special_points_frac[clean_label] = coords
            except Exception as e:
                print(f"  Warning: Could not determine special points: {e}")
                special_points_frac = {}

            # Convert all q-points to Cartesian
            q_cart = np.dot(q_frac, rec_cell)
            
        except Exception as e:
            print(f"Error generating BZ/Geometry from {poscar_path}: {e}")
            q_cart = q_frac # Fallback

    # 3. Colorscale Logic
    g = points_g
    
    # Auto-ranging logic (active determination from data)
    if gmin is None:
        q1, q3 = np.percentile(g, [25, 75])
        iqr = q3 - q1
        f_factor = filter_outliers if filter_outliers is not None else 3.0
        auto_vmin = q1 - f_factor * iqr
        # Use actual min if it's within reasonable bounds, otherwise clamp
        real_vmin = max(g.min(), auto_vmin)
    else:
        real_vmin = gmin

    if gmax is None:
        q1, q3 = np.percentile(g, [25, 75])
        iqr = q3 - q1
        f_factor = filter_outliers if filter_outliers is not None else 3.0
        auto_vmax = q3 + f_factor * iqr
        real_vmax = min(g.max(), auto_vmax)
    else:
        real_vmax = gmax

    # 4. Create Plot
    fig = go.Figure()

    # Add BZ Wireframe
    for start, end in wireframe:
        fig.add_trace(go.Scatter3d(
            x=[start[0], end[0]], y=[start[1], end[1]], z=[start[2], end[2]],
            mode='lines',
            line=dict(color='black', width=2),
            showlegend=False,
            hoverinfo='skip'
        ))

    # Add Q-points
    cb_title = "Grüneisen<br>Parameter"
    if mode == "mean": cb_title += "<br>(Mean)"
    if mode == "min": cb_title += "<br>(Min)"

    # Hover text for regular points: g and fractional q
    q_hover_text = []
    for i in range(len(points_g)):
        qf = q_frac[i]
        q_hover_text.append(f"q: [{qf[0]:.3f}, {qf[1]:.3f}, {qf[2]:.3f}]<br>g: {points_g[i]:.3f}")

    fig.add_trace(go.Scatter3d(
        x=q_cart[:, 0], y=q_cart[:, 1], z=q_cart[:, 2],
        mode='markers',
        marker=dict(
            size=4,
            color=points_g,
            colorscale='Plasma',
            cmin=real_vmin,
            cmax=real_vmax,
            colorbar=dict(title=cb_title),
            opacity=0.8
        ),
        text=q_hover_text,
        hoverinfo='text',
        name='Q-points',
        showlegend=False
    ))
    
    # Add Special Points
    if special_points_cart:
        sp_x, sp_y, sp_z, sp_text, sp_hover = [], [], [], [], []
        
        # Build KDTree for nearest neighbor search if points exist
        tree = KDTree(q_cart) if len(q_cart) > 0 else None

        for label, coord in special_points_cart.items():
            sp_x.append(coord[0])
            sp_y.append(coord[1])
            sp_z.append(coord[2])
            sp_text.append(label)
            
            qf = special_points_frac.get(label, [0,0,0])
            qf_str = f"[{qf[0]:.3f}, {qf[1]:.3f}, {qf[2]:.3f}]"

            if tree:
                dist, idx = tree.query(coord)
                val = points_g[idx]
                sp_hover.append(f"Special Point: {label}<br>q: {qf_str}<br>g: {val:.3f}")
            else:
                sp_hover.append(f"Special Point: {label}<br>q: {qf_str}")
            
        fig.add_trace(go.Scatter3d(
            x=sp_x, y=sp_y, z=sp_z,
            mode='text+markers',
            marker=dict(size=5, color='black'),
            text=sp_text,
            hovertext=sp_hover,
            hoverinfo='text',
            textposition="top center",
            textfont=dict(size=15, color="black", family="Arial Black"),
            showlegend=False,
            name='Special Points'
        ))
    elif not poscar_path:
        # Fallback if no POSCAR found (only Gamma at 0)
        # Find g at Gamma if possible
        tree = KDTree(q_cart) if len(q_cart) > 0 else None
        g_gamma_str = ""
        if tree:
            dist, idx = tree.query([0,0,0])
            g_gamma_str = f"<br>g: {points_g[idx]:.3f}"

        fig.add_trace(go.Scatter3d(
            x=[0], y=[0], z=[0],
            mode='text+markers',
            marker=dict(size=5, color='black'),
            text=['Γ'],
            hovertext=[f"Special Point: Γ<br>q: [0.000, 0.000, 0.000]{g_gamma_str}"],
            hoverinfo='text',
            textposition="top center",
            textfont=dict(size=15, color="black"),
            showlegend=False
        ))

    # Add Reciprocal Vectors (b1, b2, b3)
    if poscar_path:
        b_colors = ['red', 'green', 'blue']
        for i in range(3):
            b_vec = rec_cell[i]
            # Axis line
            fig.add_trace(go.Scatter3d(
                x=[0, b_vec[0]], y=[0, b_vec[1]], z=[0, b_vec[2]],
                mode='lines',
                line=dict(color=b_colors[i], width=4),
                name=f'b{i+1}',
                showlegend=False
            ))
            # Arrowhead
            fig.add_trace(go.Cone(
                x=[b_vec[0]], y=[b_vec[1]], z=[b_vec[2]],
                u=[b_vec[0]], v=[b_vec[1]], w=[b_vec[2]],
                sizemode="absolute", sizeref=0.1,
                colorscale=[[0, b_colors[i]], [1, b_colors[i]]],
                showscale=False,
                hoverinfo='skip',
                showlegend=False
            ))
            # Label
            fig.add_trace(go.Scatter3d(
                x=[b_vec[0] * 1.05], y=[b_vec[1] * 1.05], z=[b_vec[2] * 1.05],
                mode='text',
                text=[f"<b>b{i+1}</b>"],
                textfont=dict(color=b_colors[i], size=14),
                showlegend=False,
                hoverinfo='skip'
            ))

    mode_str = "Mean" if mode == "mean" else "Minima"
    fig.update_layout(
        title=f"3D Grüneisen Parameter Distribution<br>({mode_str} per q-point) - {Path(dat_path).stem}<br>"
              f"<span style='font-size: 12px;'>Formula: {meta_formula} | SpaceGroup: {meta_sg} | Path: {meta_path}</span>",
        scene=dict(
            xaxis_title='Qx (1/Å)',
            yaxis_title='Qy (1/Å)',
            zaxis_title='Qz (1/Å)',
            aspectmode='data',
            xaxis=dict(showbackground=True, backgroundcolor="rgb(240, 240, 240)", gridcolor="rgb(200, 200, 200)", showgrid=True),
            yaxis=dict(showbackground=True, backgroundcolor="rgb(240, 240, 240)", gridcolor="rgb(200, 200, 200)", showgrid=True),
            zaxis=dict(showbackground=True, backgroundcolor="rgb(240, 240, 240)", gridcolor="rgb(200, 200, 200)", showgrid=True),
        ),
        margin=dict(l=0, r=0, b=0, t=50)
    )
    
    fig.write_html(out_html)
    print(f"3D visualization saved to {out_html}")
    
    # Save as PDF (requires kaleido)
    out_pdf = str(Path(out_html).with_suffix(".pdf"))
    try:
        # Check if kaleido is installed effectively by trying to write
        # (Plotly usually raises ValueError or ImportError internally)
        fig.write_image(out_pdf)
        print(f"3D visualization static image saved to {out_pdf}")
    except Exception as e:
        print(f"Warning: Could not save 3D plot as PDF (requires 'pip install kaleido'): {e}")

def plot_relax_log(log_path: str, out_prefix: str = "relax_plot"):
    """Plot convergence from relax_log.txt."""
    print("plot_relax_log: To be implemented.")
