
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
    
    # Reciprocal cell (b vectors as rows, without 2pi factor from ase)
    # We multiply by 2pi to match physics convention used in plot_gruneisen_3d
    rec_cell = cell.reciprocal() * 2 * np.pi
    
    # Voronoi decomposition of the reciprocal lattice
    # 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)
    
    bz_ridges = []
    
    # Center is at 0,0,0 (index 13 in the flattened 3x3x3 grid)
    center_idx = 13 
    
    for pid, rid in zip(vor.ridge_points, vor.ridge_vertices):
        if(pid[0] == center_idx or pid[1] == center_idx):
            # rid contains indices of vertices forming the ridge (face edge)
            # We need to extract lines (pairs of vertices) from this polygon
            # Voronoi ridge vertices are ordered? Usually they form a polygon.
            # Let's connect them in order.
            
            # Note: vor.ridge_vertices can be -1 (infinity), but for Wigner-Seitz cell of lattice, it's bounded.
            if -1 in rid: continue
            
            # The rid list forms a cycle (face boundary).
            # Connect v[i] to v[i+1]
            for i in range(len(rid)):
                v1 = vor.vertices[rid[i]]
                v2 = vor.vertices[rid[(i+1)%len(rid)]]
                bz_ridges.append((v1, v2))
            
    return bz_ridges

def apply_hull_filter(q_cart, values, q_frac, special_points, tolerance=0.05):
    """
    Filter points to keep only those within the Convex Hull of special points.
    Applies a tolerance by artificially expanding the hull.
    """
    if not special_points or len(special_points) < 4:
        print("Warning: Not enough special points for Convex Hull filtering. Skipping.")
        return q_cart, values, q_frac

    try:
        from scipy.spatial import Delaunay
        
        # To apply tolerance, we expand the hull points slightly outward from the center (Gamma).
        pts = np.array(list(special_points.values()))
        # Assume Gamma (0,0,0) is center.
        expanded_pts = pts * (1.0 + tolerance)
        
        hull = Delaunay(expanded_pts)
        
        # Check if points are inside
        # find_simplex returns -1 if point is outside
        inside_mask = hull.find_simplex(q_cart) >= 0
        
        return q_cart[inside_mask], values[inside_mask], q_frac[inside_mask]
        
    except Exception as e:
        print(f"Error during Convex Hull filtering: {e}")
        return q_cart, values, q_frac

def plot_gruneisen_3d(dat_path, poscar_path=None, out_html="gruneisen_3d.html", gmin=None, gmax=None, filter_outliers=3.0, mode="mean", use_hull_filter=True, hull_tolerance=0.05, expand_symmetry=True, half_bz_fold=True):
    """
    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"
    meta_path = "N/A"
    meta_path = "N/A"
    meta_path = "N/A"
    wireframe = []
    special_points_cart = {}
    special_points_frac = {}
    rec_cell = np.eye(3)
    # 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
    q_cart = q_frac.copy()
    if mode == "min":
        points_g = q_groups['min'].values
    else:
        points_g = q_groups['mean'].values

    def _best_half_bz_axis(points_frac):
        if not points_frac:
            return 0
        pts = np.array(points_frac, dtype=float)
        if pts.ndim != 2 or pts.shape[1] != 3:
            return 0
        counts = []
        for i in range(3):
            counts.append(int(np.sum(pts[:, i] > -0.5)))
        return int(np.argmax(counts))
    
    # Half-BZ folding (auto axis to include most special points)
    if half_bz_fold and not expand_symmetry:
        axis_idx = _best_half_bz_axis(special_points_frac_vals)
        mask = q_frac[:, axis_idx] <= -0.5
        q_frac[mask] *= -1

    # 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) ---
                        if use_hull_filter:
                            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.")
                        else:
                            print("  Skipping Convex Hull Filtering (Full BZ visualization).")
    
                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. Symmetry Expansion & Filtering
        if poscar_path:
                try:
                    from ase.io import read
                    structure = read(str(poscar_path))
                    cell_tuple = (structure.get_cell(), structure.get_scaled_positions(), structure.get_atomic_numbers())
                    
                    # Reciprocal Cell
                    rec_cell = structure.get_cell().reciprocal() * 2 * np.pi
                    
                    # --- Special Points (Seekpath) ---
                    try:
                        try:
                            import seekpath
                        except ImportError:
                            import macer.externals.seekpath_bundled as seekpath
                        
                        res = seekpath.get_explicit_k_path(cell_tuple)
                        
                        # Path metadata
                        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])
        
                        # Special points
                        special_points_cart = {}
                        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}")
        
                    # --- Symmetry Expansion (Star of k) ---
                    try:
                        # Heuristic check: If q-points are many, assume Full Mesh and skip expansion
                        is_likely_full_mesh = len(q_frac) > 2000
                        
                        if expand_symmetry and not is_likely_full_mesh:
                            try:
                                import spglib
                            except ImportError:
                                import macer.externals.spglib_bundled as spglib
                            
                            dataset = spglib.get_symmetry(cell_tuple)
                            rotations = dataset['rotations']
                            
                            # Expand q-points
                            expanded_q = []
                            expanded_g = []
                            for q, g_val in zip(q_frac, points_g):
                                qs_star = np.matmul(rotations, q)
                                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)
                                print(f"  Symmetry Expansion: {len(q_groups)} -> {len(q_frac)} points.")
                                
                                # Map to 1st BZ (closest to Gamma) - Only if expanded
                                q_cart_base = np.dot(q_frac, rec_cell)
                                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)
                                
                                candidates = q_cart_base[:, np.newaxis, :] + shift_cart[np.newaxis, :, :]
                                dists = np.linalg.norm(candidates, axis=2)
                                best_indices = np.argmin(dists, axis=1)
                                best_shifts = shifts[best_indices]
                                q_frac = q_frac + best_shifts
        
                    except ImportError:
                        pass # spglib not found
                    
                    if is_likely_full_mesh:
                        print(f"  Input seems to be Full Mesh ({len(q_frac)} points). Skipping Symmetry Expansion & BZ Mapping.")
        
                    # --- Half-BZ Folding (auto axis to include most special points) ---
                    if half_bz_fold:
                        axis_idx = _best_half_bz_axis(list(special_points_frac.values()))
                        mask = q_frac[:, axis_idx] <= -0.5
                        q_frac[mask] *= -1
                    
                    q_cart = np.dot(q_frac, rec_cell)
                    
                    # --- Convex Hull Filtering ---
                    if use_hull_filter:
                        print(f"  Filtering q-points by Convex Hull (tol={hull_tolerance})...")
                        prev_len = len(q_cart)
                        q_cart, points_g, q_frac = apply_hull_filter(q_cart, points_g, q_frac, special_points_cart, hull_tolerance)
                        print(f"  Filter result: {prev_len} -> {len(q_cart)} points.")
                    else:
                        print("  Skipping Convex Hull Filtering (Full BZ visualization).")
        
                except Exception as e:
                    print(f"  Warning: Symmetry expansion/filtering failed: {e}")
                    import traceback
                    traceback.print_exc()
                    q_cart = np.dot(q_frac, rec_cell)
        
                # BZ Wireframe
                wireframe = _get_bz_wireframe(structure.get_cell())
    g = points_g
    # Check for NaNs/Infs
    valid_mask = np.isfinite(g)
    if not np.all(valid_mask):
        print(f"  Warning: Found {np.sum(~valid_mask)} invalid (NaN/Inf) Grüneisen values. They will be hidden.")
        q_cart = q_cart[valid_mask]
        q_frac = q_frac[valid_mask]
        points_g = points_g[valid_mask]
        g = points_g

    if len(g) == 0:
        print("  Error: No valid data points to plot.")
        return

    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
        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

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

    # 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=4),
            showlegend=False,
            hoverinfo='skip'
        ))

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

    # Hover text with q-point info
    q_hover_text = [f"q: [{q[0]:.3f}, {q[1]:.3f}, {q[2]:.3f}]<br>g: {v:.3f}" for q, v in zip(q_frac, points_g)]

    # Colorscale (handle 0 crossing)
    if real_vmin < 0 < real_vmax:
        midpoint = (0 - real_vmin) / (real_vmax - real_vmin)
        colorscale = [[0.0, 'rgb(0, 0, 255)'], [midpoint, 'rgb(255, 255, 255)'], [1.0, 'rgb(255, 0, 0)']]
    else:
        colorscale = 'Plasma'

    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=colorscale,
            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
    ))
    
    # Special Points
    if special_points_cart:
        sp_x = [c[0] for c in special_points_cart.values()]
        sp_y = [c[1] for c in special_points_cart.values()]
        sp_z = [c[2] for c in special_points_cart.values()]
        sp_text = list(special_points_cart.keys())
        
        # Add hover info for special points (fractional coords)
        sp_hover = []
        for label in special_points_cart.keys():
            qf = special_points_frac.get(label, [0,0,0])
            sp_hover.append(f"{label}<br>q: [{qf[0]:.3f}, {qf[1]:.3f}, {qf[2]:.3f}]")

        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'
        ))

    # Reciprocal Vectors
    if poscar_path:
        b_colors = ['red', 'green', 'blue']
        for i in range(3):
            b_vec = rec_cell[i]
            # Axis line (Mochi style: simple line)
            fig.add_trace(go.Scatter3d(
                x=[0, b_vec[0]], y=[0, b_vec[1]], z=[0, b_vec[2]],
                mode='lines+text',
                line=dict(color=b_colors[i], width=6),
                text=['', f"<b>b{i+1}</b>"],
                textfont=dict(color=b_colors[i], size=20),
                hoverinfo='skip',
                showlegend=False
            ))

    # Layout & Save
    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(visible=False),
            yaxis=dict(visible=False),
            zaxis=dict(visible=False),
        ),
        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
    out_pdf = str(Path(out_html).with_suffix(".pdf"))
    try:
        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_gruneisen_3d_style(
    dat_path,
    poscar_path="POSCAR",
    out_prefix="average_gruneisen_BZ",
    cmin=-5,
    cmax=5,
    cut_f=0.0,
    original_position=False,
    fill_bz=False,
    plot_rc=False,
    expand_symmetry=True,
    full_bz=False,
    half_bz_axis="auto",
    use_hull_filter=False,
    hull_tolerance=0.05,
    special_points=True,
    time_reversal=True,
    symprec=1e-2,
    mode="mean",
):
    """Plot average Gruneisen parameter in 3D BZ (style variant)."""
    meta_formula = "N/A"
    meta_sg = "N/A"
    try:
        import plotly.graph_objs as go
        import plotly.offline as pyo
    except ImportError:
        print("Error: plotly is required for this plot.")
        return
    meta_formula = "N/A"
    meta_sg = "N/A"

    meta_formula = "N/A"
    meta_sg = "N/A"

    # Metadata from dat header (if present)
    meta_formula = "N/A"
    meta_sg = "N/A"
    meta_path = "N/A"
    try:
        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
    except Exception:
        pass

    def get_lattice_vector(poscar_path_in):
        with open(poscar_path_in, "r") as f:
            f.readline()
            scale = float(f.readline())
            lattice_vectors = []
            for _ in range(3):
                vector = f.readline().split()
                lattice_vectors.append([float(component) for component in vector])
        return np.array(lattice_vectors) * scale

    def get_brillouin_zone_3d(cell):
        cell = np.asarray(cell, dtype=float)
        assert cell.shape == (3, 3)
        px, py, pz = np.tensordot(cell, np.mgrid[-1:2, -1:2, -1:2], axes=[0, 0])
        points = np.c_[px.ravel(), py.ravel(), pz.ravel()]
        from scipy.spatial import Voronoi
        vor = Voronoi(points)
        bz_ridges = []
        bz_vertices = []
        for pid, rid in zip(vor.ridge_points, vor.ridge_vertices):
            if pid[0] == 13 or pid[1] == 13:
                bz_ridges.append(vor.vertices[np.r_[rid, [rid[0]]]])
                bz_vertices += rid
        bz_vertices = list(set(bz_vertices))
        return vor.vertices[bz_vertices], bz_ridges, None

    def q_position_fixer(q_position, fix_position):
        if fix_position and q_position[2] <= -0.5:
            q_position = q_position * (-1)
        return q_position

    def get_gruneisen_from_dat(dat_path_in, cut_f_in, fix_position_in, mode_in):
        from collections import OrderedDict
        q_to_gamma = OrderedDict()
        with open(dat_path_in, "r") as f:
            for line in f:
                if not line.strip() or line.lstrip().startswith("#"):
                    continue
                parts = line.split()
                if len(parts) < 7:
                    continue
                qx, qy, qz = map(float, parts[0:3])
                frequency = float(parts[5])
                gruneisen = float(parts[6])
                q_position = q_position_fixer(np.array([qx, qy, qz], dtype=float), fix_position_in)
                key = tuple(q_position.tolist())
                if key not in q_to_gamma:
                    q_to_gamma[key] = []
                if frequency > cut_f_in:
                    q_to_gamma[key].append(gruneisen)

        q_positions = []
        g_values = []
        for q_position, gamma_list in q_to_gamma.items():
            if not gamma_list:
                continue
            q_position = np.array(q_position, dtype=float)
            if mode_in == "min":
                gruneisen_average = np.min(gamma_list)
            else:
                gruneisen_average = np.mean(gamma_list)
            q_positions.append(q_position)
            g_values.append(gruneisen_average)
        return np.array(q_positions, dtype=float), np.array(g_values, dtype=float)

    def expand_qpoints_full_bz(q_frac, g_values, poscar_path_in, rec_cell_in, symprec_in, time_reversal_in):
        try:
            from ase.io import read as ase_read
        except Exception as e:
            print(f"Warning: ASE not available for symmetry expansion: {e}")
            return q_frac, g_values
        try:
            try:
                import spglib
            except Exception:
                import macer.externals.spglib_bundled as spglib
        except Exception as e:
            print(f"Warning: spglib not available for symmetry expansion: {e}")
            return q_frac, g_values
        try:
            structure = ase_read(str(poscar_path_in))
            cell_tuple = (structure.get_cell(), structure.get_scaled_positions(), structure.get_atomic_numbers())
            rotations = None
            for symp in (symprec_in, 1e-5, 1e-4, 1e-3, 1e-2):
                dataset = spglib.get_symmetry(cell_tuple, symprec=symp)
                rotations = dataset.get("rotations") if dataset else None
                if rotations is not None and len(rotations) > 1:
                    break
            if rotations is None:
                raise ValueError("spglib did not return symmetry rotations.")
            if len(rotations) == 1:
                print("Warning: Only identity symmetry found. Skipping symmetry expansion.")
                return q_frac, g_values
        except Exception as e:
            print(f"Warning: Failed to load symmetry from POSCAR: {e}")
            return q_frac, g_values

        expanded_q = []
        expanded_g = []
        for q, g_val in zip(q_frac, g_values):
            qs_star = np.matmul(rotations, q)
            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 not expanded_q:
            return q_frac, g_values
        q_frac = np.vstack(expanded_q)
        g_values = np.concatenate(expanded_g)

        if time_reversal_in:
            q_tr = -q_frac
            g_tr = g_values.copy()
            q_frac = np.vstack([q_frac, q_tr])
            g_values = np.concatenate([g_values, g_tr])

        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_in)
        q_cart_base = np.dot(q_frac, rec_cell_in)
        candidates = q_cart_base[:, np.newaxis, :] + shift_cart[np.newaxis, :, :]
        dists = np.linalg.norm(candidates, axis=2)
        best_indices = np.argmin(dists, axis=1)
        best_shifts = shifts[best_indices]
        q_frac = q_frac + best_shifts
        # unique-ify after shifts
        qg = np.c_[np.round(q_frac, 8), np.round(g_values, 8)]
        _, idx = np.unique(qg[:, :3], axis=0, return_index=True)
        q_frac = q_frac[idx]
        g_values = g_values[idx]
        return q_frac, g_values

    fix_position = not original_position
    meta_formula = "N/A"
    meta_sg = "N/A"
    from ase.io import read as ase_read
    try:
        structure = ase_read(str(poscar_path))
        meta_formula = structure.get_chemical_formula()
        try:
            import spglib
        except Exception:
            import macer.externals.spglib_bundled as spglib
        cell_tuple = (structure.get_cell(), structure.get_scaled_positions(), structure.get_atomic_numbers())
        try:
            meta_sg = spglib.get_spacegroup(cell_tuple, symprec=1e-5)
        except Exception:
            meta_sg = "N/A"
        try:
            try:
                import seekpath
            except Exception:
                import macer.externals.seekpath_bundled as seekpath
            res = seekpath.get_explicit_k_path(cell_tuple)
            if "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])
        except Exception:
            meta_path = "N/A"
    except Exception:
        meta_formula = "N/A"
        meta_sg = "N/A"
        meta_path = "N/A"

    lat = get_lattice_vector(poscar_path)
    rec_cell = np.linalg.inv(lat).T

    def _get_special_points(poscar_path_in, rec_cell_in):
        try:
            try:
                import seekpath
            except Exception:
                import macer.externals.seekpath_bundled as seekpath
            from ase.io import read
            structure = read(str(poscar_path_in))
            cell_tuple = (structure.get_cell(), structure.get_scaled_positions(), structure.get_atomic_numbers())
            res = seekpath.get_explicit_k_path(cell_tuple)
            sp_frac = {}
            sp_cart = {}
            if "point_coords" in res:
                for label, coords in res["point_coords"].items():
                    cart_coord = np.dot(coords, rec_cell_in)
                    clean_label = label.replace("GAMMA", "Γ").replace("GM", "Γ")
                    sp_cart[clean_label] = cart_coord
                    sp_frac[clean_label] = coords
            return sp_frac, sp_cart
        except Exception:
            return {}, {}

    def _best_half_bz_axis(points_frac):
        if not points_frac:
            return 0
        pts = np.array(list(points_frac.values()), dtype=float)
        if pts.ndim != 2 or pts.shape[1] != 3:
            return 0
        counts = []
        for i in range(3):
            counts.append(int(np.sum(pts[:, i] >= 0.0)))
        return int(np.argmax(counts))

    if expand_symmetry:
        fix_position = False

    q_positions, g_values = get_gruneisen_from_dat(dat_path, cut_f, fix_position, mode)
    if full_bz:
        uq = len(np.unique(np.round(q_positions, 8), axis=0))
        if uq < 200:
            print("Warning: q-point set is very small; input may be a band path, not a full mesh. Full BZ fill may be incomplete.")
    if expand_symmetry:
        q_positions, g_values = expand_qpoints_full_bz(q_positions, g_values, poscar_path, rec_cell, symprec, time_reversal)

    q_positions_all = q_positions.copy()
    g_values_all = g_values.copy()

    sp_frac, sp_cart = _get_special_points(poscar_path, rec_cell)
    if half_bz_axis == "full":
        full_bz = True

    if not full_bz:
        if half_bz_axis == "b1":
            axis_idx = 0
        elif half_bz_axis == "b2":
            axis_idx = 1
        elif half_bz_axis == "b3":
            axis_idx = 2
        else:
            axis_idx = _best_half_bz_axis(sp_frac)
        mask = q_positions[:, axis_idx] >= 0.0
        q_positions = q_positions[mask]
        g_values = g_values[mask]

    q_positions_ac = np.dot(q_positions, rec_cell)
    if use_hull_filter:
        q_positions_ac, g_values, q_positions = apply_hull_filter(
            q_positions_ac,
            g_values,
            q_positions,
            sp_cart,
            tolerance=hull_tolerance
        )
    data_array = np.column_stack([q_positions, q_positions_ac, g_values])
    if fill_bz:
        add_data = data_array.copy()
        add_data[:, list(range(0, 6))] *= (-1)
        data_array = np.concatenate([data_array, add_data])

    out_csv = f"{out_prefix}.csv"
    with open(out_csv, "w") as fp:
        fp.write("qx_rc,qy_rc,qz_rc,qx_ac,qy_ac,qz_ac,gruneisen_average\n")
        np.savetxt(fp, data_array, fmt="%f", delimiter=",")

    df_plot = pd.read_csv(out_csv)

    def axis_setting(label, show_ticks, rng):
        return dict(
            showticklabels=show_ticks,
            title=dict(text=label, font=dict(family="Arial", size=30, color="black")),
            tickfont=dict(family="Arial", size=16, color="black"),
            showbackground=True,
            backgroundcolor="rgba(210,210,210,0.9)",
            showgrid=True,
            gridcolor="rgba(255,255,255,0.9)",
            gridwidth=2.5,
            zeroline=False,
            range=rng,
        )

    if plot_rc:
        qx = df_plot["qx_rc"]
        qy = df_plot["qy_rc"]
        qz = df_plot["qz_rc"]
    else:
        qx = df_plot["qx_ac"]
        qy = df_plot["qy_ac"]
        qz = df_plot["qz_ac"]

    if full_bz:
        pad_x = 0.25 * (float(np.max(qx)) - float(np.min(qx)) + 1e-12)
        pad_y = 0.25 * (float(np.max(qy)) - float(np.min(qy)) + 1e-12)
        pad_z = 0.25 * (float(np.max(qz)) - float(np.min(qz)) + 1e-12)
        xrng = [float(np.min(qx)) - pad_x, float(np.max(qx)) + pad_x]
        yrng = [float(np.min(qy)) - pad_y, float(np.max(qy)) + pad_y]
        zrng = [float(np.min(qz)) - pad_z, float(np.max(qz)) + pad_z]
    else:
        max_x = float(np.max(np.abs(qx)))
        max_y = float(np.max(np.abs(qy)))
        max_z = float(np.max(np.abs(qz)))
        if not plot_rc:
            max_x = max(max_x, float(np.max(np.abs(rec_cell[:, 0]))))
            max_y = max(max_y, float(np.max(np.abs(rec_cell[:, 1]))))
            max_z = max(max_z, float(np.max(np.abs(rec_cell[:, 2]))))
        pad_x = 0.25 * (max_x + 1e-12)
        pad_y = 0.25 * (max_y + 1e-12)
        pad_z = 0.25 * (max_z + 1e-12)
        xrng = [-max_x - pad_x, max_x + pad_x]
        yrng = [-max_y - pad_y, max_y + pad_y]
        zrng = [-max_z - pad_z, max_z + pad_z]

    colorbar_layout = dict(
        title=dict(text="&#947;<sub>q</sub>", font=dict(family="Times New Roman", size=35, color="black")),
        tickfont=dict(family="Arial", size=24, color="black"),
        len=0.9,
        ticks="outside",
        tickwidth=2,
    )
    cool_colorscale = [
        [0.0, "#3b4cc0"],
        [0.1111111111111111, "#5977e3"],
        [0.2222222222222222, "#7aa6f2"],
        [0.3333333333333333, "#a9cfe5"],
        [0.4444444444444444, "#d4e4f1"],
        [0.5555555555555556, "#f2d9c2"],
        [0.6666666666666666, "#f3a481"],
        [0.7777777777777778, "#e36a53"],
        [0.8888888888888888, "#c93b3b"],
        [1.0, "#b40426"],
    ]
    marker_setting = dict(
        size=5,
        color=df_plot["gruneisen_average"],
        colorscale=cool_colorscale,
        cmin=cmin,
        cmax=cmax,
        opacity=1.0,
        colorbar=colorbar_layout,
    )
    trace = go.Scatter3d(
        x=qx,
        y=qy,
        z=qz,
        customdata=np.stack((df_plot["qx_rc"], df_plot["qy_rc"], df_plot["qz_rc"], df_plot["gruneisen_average"]), axis=-1),
        mode="markers",
        hovertemplate="q-position: (%{customdata[0]}, %{customdata[1]}, %{customdata[2]})<br>&#947;<sub>q</sub> average: %{customdata[3]}<extra></extra>",
        marker=marker_setting,
    )

    layout = go.Layout(
        margin=dict(l=0, r=0, b=0, t=0),
        showlegend=False,
        scene=dict(
            xaxis=axis_setting("<i>k<sub>x</sub></i>", plot_rc, xrng),
            yaxis=axis_setting("<i>k<sub>y</sub></i>", plot_rc, yrng),
            zaxis=axis_setting("<i>k<sub>z</sub></i>", plot_rc, zrng),
        ),
    )

    fig = go.Figure(data=trace, layout=layout)

    # Outer pale planes removed to avoid extra gray box outside the grid

    if not plot_rc:
        _, edges, _ = get_brillouin_zone_3d(rec_cell)
        basis_clrs = ["red", "green", "blue"]
        basis_labs = ["<i>b<sub>1</sub></i>", "<i>b<sub>2</sub></i>", "<i>b<sub>3</sub></i>"]
        basis_text_traces = []
        basis_label_pos = []
        outer = max(abs(xrng[0]), abs(xrng[1]), abs(yrng[0]), abs(yrng[1]), abs(zrng[0]), abs(zrng[1]))
        for ii, basis in enumerate(rec_cell):
            bx, by, bz = basis
            vec = np.array([bx, by, bz], dtype=float)
            vnorm = np.linalg.norm(vec)
            if vnorm == 0:
                continue
            unit = vec / vnorm
            arrow_end = unit * (0.95 * outer)
            label_pos = unit * (1.02 * outer)
            fig.add_trace(go.Scatter3d(
                x=[0, arrow_end[0]], y=[0, arrow_end[1]], z=[0, arrow_end[2]],
                opacity=0.9, hoverinfo="skip", mode="lines",
                line=dict(color=basis_clrs[ii], width=6),
            ))
            fig.add_trace(go.Cone(
                x=[arrow_end[0]], y=[arrow_end[1]], z=[arrow_end[2]],
                u=[unit[0]], v=[unit[1]], w=[unit[2]],
                sizemode="absolute",
                sizeref=0.12 * outer,
                anchor="tip",
                showscale=False,
                colorscale=[[0, basis_clrs[ii]], [1, basis_clrs[ii]]],
                hoverinfo="skip",
                opacity=0.95,
            ))
            basis_label_pos.append(label_pos)
            basis_text_traces.append(go.Scatter3d(
                x=[label_pos[0]], y=[label_pos[1]], z=[label_pos[2]],
                mode="text",
                text=[basis_labs[ii]],
                textfont=dict(color=basis_clrs[ii], size=20),
                hoverinfo="skip",
                opacity=1.0,
                showlegend=False,
            ))
        # Expand axis ranges to include basis labels outside the plot box
        if basis_label_pos:
            basis_label_pos = np.array(basis_label_pos)
            xrng = [min(xrng[0], float(basis_label_pos[:, 0].min())),
                    max(xrng[1], float(basis_label_pos[:, 0].max()))]
            yrng = [min(yrng[0], float(basis_label_pos[:, 1].min())),
                    max(yrng[1], float(basis_label_pos[:, 1].max()))]
            zrng = [min(zrng[0], float(basis_label_pos[:, 2].min())),
                    max(zrng[1], float(basis_label_pos[:, 2].max()))]
            fig.update_layout(scene=dict(
                xaxis=axis_setting("<i>k<sub>x</sub></i>", plot_rc, xrng),
                yaxis=axis_setting("<i>k<sub>y</sub></i>", plot_rc, yrng),
                zaxis=axis_setting("<i>k<sub>z</sub></i>", plot_rc, zrng),
            ))
        for edge in edges:
            fig.add_trace(go.Scatter3d(
                x=edge[:, 0], y=edge[:, 1], z=edge[:, 2],
                opacity=0.8, hoverinfo="skip", mode="lines",
                line=dict(color="black", width=5),
            ))
        for tr in basis_text_traces:
            fig.add_trace(tr)

        if special_points:
            try:
                try:
                    import seekpath
                except Exception:
                    import macer.externals.seekpath_bundled as seekpath
                try:
                    from scipy.spatial import KDTree
                except Exception:
                    KDTree = None
                structure = read(str(poscar_path))
                cell_tuple = (structure.get_cell(), structure.get_scaled_positions(), structure.get_atomic_numbers())
                res = seekpath.get_explicit_k_path(cell_tuple)
                if "point_coords" in res:
                    sp_cart = []
                    sp_labels = []
                    sp_hover = []
                    if KDTree is not None and len(q_positions_all) > 0:
                        tree = KDTree(q_positions_all)
                    else:
                        tree = None
                    for label, coords in res["point_coords"].items():
                        cart_coord = np.dot(coords, rec_cell)
                        clean_label = label.replace("GAMMA", "Γ").replace("GM", "Γ")
                        sp_cart.append(cart_coord)
                        sp_labels.append(clean_label)
                        if tree is not None:
                            _, idx = tree.query(coords)
                            g_val = g_values_all[idx]
                            sp_hover.append(
                                f"{clean_label}<br>q: [{coords[0]:.3f}, {coords[1]:.3f}, {coords[2]:.3f}]<br>"
                                f"γ: {g_val:.3f}"
                            )
                        else:
                            sp_hover.append(
                                f"{clean_label}<br>q: [{coords[0]:.3f}, {coords[1]:.3f}, {coords[2]:.3f}]"
                            )
                    if sp_cart:
                        sp_cart = np.array(sp_cart)
                        fig.add_trace(go.Scatter3d(
                            x=sp_cart[:, 0], y=sp_cart[:, 1], z=sp_cart[:, 2],
                            mode="markers+text",
                            text=sp_labels,
                            textposition="top center",
                            textfont=dict(color="black", size=20),
                            marker=dict(color="black", size=4),
                            hovertext=sp_hover,
                            hoverinfo="text",
                            showlegend=False,
                        ))
            except Exception as e:
                print(f"Warning: Special q-point labeling skipped: {e}")

    mode_str = "Gruneisen Full (Average)" if mode == "mean" else "Gruneisen Minima"
    title_style = "font-size: 12px;"
    fig.update_layout(
        title=dict(
            text=f"<span style='{title_style}'>{mode_str}</span><br>"
                 f"<span style='{title_style}'>Formula: {meta_formula} | SpaceGroup: {meta_sg} | Path: {meta_path}</span>",
            x=0.5,
            xanchor="center"
        ),
        margin=dict(l=0, r=0, b=0, t=60)
    )
    out_html = f"{out_prefix}.html"
    pyo.plot(fig, filename=out_html, auto_open=True)
    out_pdf = f"{out_prefix}.pdf"
    try:
        fig.write_image(out_pdf)
    except Exception as e:
        print(f"Warning: PDF export failed (requires 'kaleido'). {e}")

def plot_bz_3d(
    poscar_path,
    out_html="BZ.html",
    plot_rc=False,
    special_points=True,
    symprec=1e-5,
):
    """Plot 1BZ wireframe and special q-points."""
    meta_formula = "N/A"
    meta_sg = "N/A"
    meta_path = "N/A"
    try:
        import plotly.graph_objs as go
        import plotly.offline as pyo
    except ImportError:
        print("Error: plotly is required for this plot.")
        return

    def get_lattice_vector(poscar_path_in):
        with open(poscar_path_in, "r") as f:
            f.readline()
            scale = float(f.readline())
            lattice_vectors = []
            for _ in range(3):
                vector = f.readline().split()
                lattice_vectors.append([float(component) for component in vector])
        return np.array(lattice_vectors) * scale

    def get_brillouin_zone_3d(cell):
        cell = np.asarray(cell, dtype=float)
        assert cell.shape == (3, 3)
        px, py, pz = np.tensordot(cell, np.mgrid[-1:2, -1:2, -1:2], axes=[0, 0])
        points = np.c_[px.ravel(), py.ravel(), pz.ravel()]
        from scipy.spatial import Voronoi
        vor = Voronoi(points)
        bz_ridges = []
        bz_vertices = []
        for pid, rid in zip(vor.ridge_points, vor.ridge_vertices):
            if pid[0] == 13 or pid[1] == 13:
                bz_ridges.append(vor.vertices[np.r_[rid, [rid[0]]]])
                bz_vertices += rid
        bz_vertices = list(set(bz_vertices))
        return vor.vertices[bz_vertices], bz_ridges

    # Metadata from POSCAR (formula, spacegroup, k-path)
    try:
        from ase.io import read as ase_read
        structure = ase_read(str(poscar_path))
        meta_formula = structure.get_chemical_formula()
        try:
            import spglib
        except Exception:
            import macer.externals.spglib_bundled as spglib
        cell_tuple = (structure.get_cell(), structure.get_scaled_positions(), structure.get_atomic_numbers())
        try:
            meta_sg = spglib.get_spacegroup(cell_tuple, symprec=symprec)
        except Exception:
            meta_sg = "N/A"
        try:
            try:
                import seekpath
            except Exception:
                import macer.externals.seekpath_bundled as seekpath
            for sp in [symprec, 1e-4, 1e-3]:
                res = seekpath.get_explicit_k_path(cell_tuple, symprec=sp)
                if "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])
                    break
        except Exception as e:
            meta_path = "N/A"
            print(f"Warning: seekpath failed to determine k-path (symprec={symprec}): {e}")
    except Exception:
        meta_formula = "N/A"
        meta_sg = "N/A"
        meta_path = "N/A"

    lat = get_lattice_vector(poscar_path)
    rec_cell = np.linalg.inv(lat).T
    v_bz, edges = get_brillouin_zone_3d(rec_cell)
    if plot_rc:
        v_plot = np.dot(v_bz, np.linalg.inv(rec_cell))
    else:
        v_plot = v_bz

    extent = np.max(np.abs(v_plot), axis=0)
    pad = 0.25 * (extent + 1e-12)
    base_xrng = [-extent[0] - pad[0], extent[0] + pad[0]]
    base_yrng = [-extent[1] - pad[1], extent[1] + pad[1]]
    base_zrng = [-extent[2] - pad[2], extent[2] + pad[2]]
    # Display ranges will expand later to include basis labels outside planes
    xrng = list(base_xrng)
    yrng = list(base_yrng)
    zrng = list(base_zrng)

    def axis_setting(label, show_ticks, rng):
        return dict(
            showticklabels=show_ticks,
            title=dict(text=label, font=dict(family="Arial", size=30, color="black")),
            tickfont=dict(family="Arial", size=16, color="black"),
            showbackground=True,
            backgroundcolor="rgba(210,210,210,0.9)",
            showgrid=True,
            gridcolor="rgba(255,255,255,0.9)",
            gridwidth=2.5,
            zeroline=False,
            range=rng,
        )

    fig = go.Figure()

    # Outer pale planes removed to avoid extra gray box outside the grid
    pad_box = 0.16
    pad_x = pad_box * (base_xrng[1] - base_xrng[0])
    pad_y = pad_box * (base_yrng[1] - base_yrng[0])
    pad_z = pad_box * (base_zrng[1] - base_zrng[0])

    # BZ edges
    for edge in edges:
        if plot_rc:
            edge = np.dot(edge, np.linalg.inv(rec_cell))
        fig.add_trace(go.Scatter3d(
            x=edge[:, 0], y=edge[:, 1], z=edge[:, 2],
            opacity=0.8, hoverinfo="skip", mode="lines",
            line=dict(color="black", width=5),
        ))

    # Reciprocal basis vectors (cart only)
    if not plot_rc:
        basis_clrs = ["red", "green", "blue"]
        basis_labs = ["<i>b<sub>1</sub></i>", "<i>b<sub>2</sub></i>", "<i>b<sub>3</sub></i>"]
        basis_text_traces = []
        for ii, basis in enumerate(rec_cell):
            bx, by, bz = basis
            vec = np.array([bx, by, bz], dtype=float)
            vnorm = np.linalg.norm(vec)
            if vnorm == 0:
                continue
            unit = vec / vnorm
            plane_outer = max(
                abs(base_xrng[0] - pad_x), abs(base_xrng[1] + pad_x),
                abs(base_yrng[0] - pad_y), abs(base_yrng[1] + pad_y),
                abs(base_zrng[0] - pad_z), abs(base_zrng[1] + pad_z),
            )
            arrow_end = unit * (1.02 * plane_outer)
            label_pos = unit * (1.10 * plane_outer)
            # Arrow shaft
            fig.add_trace(go.Scatter3d(
                x=[0, arrow_end[0]], y=[0, arrow_end[1]], z=[0, arrow_end[2]],
                opacity=0.9, hoverinfo="skip", mode="lines",
                line=dict(color=basis_clrs[ii], width=6),
            ))
            # Arrow head (cone)
            cone_len = 0.16 * plane_outer
            fig.add_trace(go.Cone(
                x=[arrow_end[0]], y=[arrow_end[1]], z=[arrow_end[2]],
                u=[unit[0]], v=[unit[1]], w=[unit[2]],
                sizemode="absolute",
                sizeref=cone_len,
                anchor="tip",
                showscale=False,
                colorscale=[[0, basis_clrs[ii]], [1, basis_clrs[ii]]],
                hoverinfo="skip",
                opacity=0.95,
            ))
            # Label slightly beyond arrow tip (outside planes)
            basis_text_traces.append(go.Scatter3d(
                x=[label_pos[0]], y=[label_pos[1]], z=[label_pos[2]],
                mode="text",
                text=[basis_labs[ii]],
                textfont=dict(color=basis_clrs[ii], size=20),
                hoverinfo="skip",
                opacity=1.0,
                showlegend=False,
            ))

        # Expand display ranges to include labels outside the planes
        if basis_text_traces:
            all_x = [t.x[0] for t in basis_text_traces] + [base_xrng[0], base_xrng[1]]
            all_y = [t.y[0] for t in basis_text_traces] + [base_yrng[0], base_yrng[1]]
            all_z = [t.z[0] for t in basis_text_traces] + [base_zrng[0], base_zrng[1]]
            xrng = [min(all_x), max(all_x)]
            yrng = [min(all_y), max(all_y)]
            zrng = [min(all_z), max(all_z)]

    # Special points (seekpath)
    if special_points:
        try:
            try:
                import seekpath
            except Exception:
                import macer.externals.seekpath_bundled as seekpath
            from ase.io import read as ase_read
            structure = ase_read(str(poscar_path))
            cell_tuple = (structure.get_cell(), structure.get_scaled_positions(), structure.get_atomic_numbers())
            res = seekpath.get_explicit_k_path(cell_tuple, symprec=symprec)
            if "point_coords" in res:
                sp_cart = []
                sp_labels = []
                sp_hover = []
                for label, coords in res["point_coords"].items():
                    cart_coord = np.dot(coords, rec_cell)
                    clean_label = label.replace("GAMMA", "Γ").replace("GM", "Γ")
                    sp_cart.append(cart_coord if not plot_rc else coords)
                    sp_labels.append(clean_label)
                    sp_hover.append(
                        f"{clean_label}<br>q: [{coords[0]:.3f}, {coords[1]:.3f}, {coords[2]:.3f}]"
                    )
                if sp_cart:
                    sp_cart = np.array(sp_cart)
                    fig.add_trace(go.Scatter3d(
                        x=sp_cart[:, 0], y=sp_cart[:, 1], z=sp_cart[:, 2],
                        mode="markers+text",
                        text=sp_labels,
                        textposition="top center",
                        textfont=dict(color="black", size=20),
                        marker=dict(color="black", size=4),
                        hovertext=sp_hover,
                        hoverinfo="text",
                        showlegend=False,
                    ))
        except Exception as e:
            print(f"Warning: Special q-point labeling skipped: {e}")

    title_style = "font-size: 12px;"
    fig.update_layout(
        title=dict(
            text=f"<span style='{title_style}'>Brillouin Zone</span><br>"
                 f"<span style='{title_style}'>Formula: {meta_formula} | SpaceGroup: {meta_sg} | Path: {meta_path}</span>",
            x=0.5,
            xanchor="center"
        ),
        margin=dict(l=0, r=0, b=0, t=60),
        showlegend=False,
        scene=dict(
            xaxis=axis_setting("<i>k<sub>x</sub></i>", plot_rc, xrng),
            yaxis=axis_setting("<i>k<sub>y</sub></i>", plot_rc, yrng),
            zaxis=axis_setting("<i>k<sub>z</sub></i>", plot_rc, zrng),
            annotations=[],
        ),
    )
    for tr in basis_text_traces:
        fig.add_trace(tr)

    pyo.plot(fig, filename=out_html, auto_open=True)

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.")
