import numpy as np
import pandas as pd
import plotly.graph_objs as go
import plotly.offline as pyo
from scipy.spatial import Voronoi
import argparse
import sys
import os
from pathlib import Path

# Add project root to sys.path to access macer modules if needed
sys.path.append(str(Path(__file__).resolve().parent.parent))

def get_reciprocal_lattice(poscar_path):
    """
    Read POSCAR and return reciprocal lattice vectors (without 2pi factor).
    Returns: ndarray shape (3, 3) where rows are b1, b2, b3
    """
    try:
        with open(poscar_path, 'r') as f:
            lines = f.readlines()
            scale = float(lines[1].strip())
            if scale < 0:
                print("Warning: Negative scaling factor (volume) in POSCAR not fully supported.")
                scale = 1.0 
            
            a1 = np.array([float(x) for x in lines[2].split()])
            a2 = np.array([float(x) for x in lines[3].split()])
            a3 = np.array([float(x) for x in lines[4].split()])
            
            lat = np.array([a1, a2, a3]) * scale
            rec_lat = np.linalg.inv(lat).T
            return rec_lat, lat
    except Exception as e:
        print(f"Error reading POSCAR: {e}")
        sys.exit(1)

def get_special_points(poscar_path, rec_cell):
    """
    Get special points using seekpath.
    """
    special_points_cart = {}
    try:
        from ase.io import read
        structure = read(str(poscar_path))
        cell_tuple = (structure.get_cell(), structure.get_scaled_positions(), structure.get_atomic_numbers())
        
        try:
            import seekpath
        except ImportError:
            # Try bundled seekpath if available
            try:
                import macer.externals.seekpath_bundled as seekpath
            except ImportError:
                print("Warning: seekpath not found. Special points will not be displayed.")
                return {}

        res = seekpath.get_explicit_k_path(cell_tuple)
        
        for label, coords in res['point_coords'].items():
            # coords are fractional
            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}")
        
    return special_points_cart

def get_brillouin_zone_3d(cell):
    """
    Generate BZ using Voronoi.
    """
    cell = np.asarray(cell, dtype=float)
    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()]
    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

    return list(set(bz_vertices)), 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
        
        # hull_pts = np.array(list(special_points.values()))
        
        # To apply tolerance, we expand the hull points slightly outward from the center (Gamma).
        # Assuming Gamma is at (0,0,0) or close to the center of the set.
        # Shift to center, scale, shift back.
        pts = np.array(list(special_points.values()))
        center = np.mean(pts, axis=0)
        
        # Expand points: New = Center + (Old - Center) * (1 + tol)
        # However, for BZ which is around 0, simply scaling vectors from 0 works if 0 is center.
        # Let's assume 0 is center (Gamma).
        expanded_pts = pts * (1.0 + tolerance)
        
        # Include original points too to ensure we don't lose the shape? 
        # Actually just expanding the boundary vertices is enough to cover more volume.
        # But to be safe, let's use the expanded set.
        
        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 parse_args():
    parser = argparse.ArgumentParser(description="Visualize Grüneisen Parameter on 3D Brillouin Zone")
    parser.add_argument("-i", "--input", default="gruneisen-POSCAR_full.dat", help="Input Grüneisen .dat file")
    parser.add_argument("-p", "--poscar", default="POSCAR", help="Input POSCAR file")
    parser.add_argument("-o", "--output", default="gruneisen_bz", help="Output filename prefix (default: gruneisen_bz)")
    parser.add_argument("--gmin", type=float, default=-5, help="Colorbar min")
    parser.add_argument("--gmax", type=float, default=5, help="Colorbar max")
    parser.add_argument("--cut-f", type=float, default=0.0, help="Cutoff frequency (THz)")
    parser.add_argument("--no-fix", action="store_true", help="Disable half-BZ folding (show original)")
    parser.add_argument("--hull-filter", action="store_true", help="Apply Convex Hull filtering (IBZ-like view)")
    parser.add_argument("--hull-tol", type=float, default=0.05, help="Tolerance for hull filtering (default: 0.05)")
    parser.add_argument("--pdf", action="store_true", help="Save static PDF images (requires kaleido)")
    return parser.parse_args()

def plot_bz_with_data(q_cart, q_frac, values, rec_cell, special_points, title_suffix, output_name, args, colorscale):
    """
    Core plotting function.
    """
    # 1. Q-point Scatter
    hover_text = [f"q: [{q[0]:.3f}, {q[1]:.3f}, {q[2]:.3f}]<br>g: {v:.3f}" for q, v in zip(q_frac, values)]
    
    scatter = go.Scatter3d(
        x=q_cart[:, 0], y=q_cart[:, 1], z=q_cart[:, 2],
        mode='markers',
        marker=dict(
            size=4,
            color=values,
            colorscale=colorscale,
            cmin=args.gmin,
            cmax=args.gmax,
            colorbar=dict(title='Grüneisen'),
            opacity=0.8
        ),
        text=hover_text,
        hoverinfo='text',
        name='Q-points'
    )
    
    data_traces = [scatter]

    # 2. BZ Wireframe
    _, ridges = get_brillouin_zone_3d(rec_cell)
    for edge in ridges:
        line = go.Scatter3d(
            x=edge[:, 0], y=edge[:, 1], z=edge[:, 2],
            mode='lines',
            line=dict(color='black', width=4),
            hoverinfo='skip',
            showlegend=False
        )
        data_traces.append(line)

    # 3. Reciprocal Basis Vectors
    basis_clrs = ['red', 'green', 'blue']
    basis_labs = ['b1', 'b2', 'b3']
    for i, basis in enumerate(rec_cell):
        vec_line = go.Scatter3d(
            x=[0, basis[0]], y=[0, basis[1]], z=[0, basis[2]],
            mode='lines+text',
            line=dict(color=basis_clrs[i], width=5),
            text=['', basis_labs[i]],
            textfont=dict(color=basis_clrs[i], size=16),
            hoverinfo='skip',
            showlegend=False
        )
        data_traces.append(vec_line)

    # 4. Special Points
    if special_points:
        sp_x = [c[0] for c in special_points.values()]
        sp_y = [c[1] for c in special_points.values()]
        sp_z = [c[2] for c in special_points.values()]
        sp_text = list(special_points.keys())
        
        sp_trace = go.Scatter3d(
            x=sp_x, y=sp_y, z=sp_z,
            mode='text+markers',
            marker=dict(size=5, color='black'),
            text=sp_text,
            textposition="top center",
            textfont=dict(size=14, color="black", family="Arial Black"),
            hoverinfo='text',
            showlegend=False
        )
        data_traces.append(sp_trace)

    # Layout
    layout = go.Layout(
        title=f"Grüneisen Parameter on BZ ({title_suffix})",
        scene=dict(
            xaxis=dict(title='kx', visible=False),
            yaxis=dict(title='ky', visible=False),
            zaxis=dict(title='kz', visible=False),
            aspectmode='data'
        ),
        margin=dict(l=0, r=0, b=0, t=50)
    )

    fig = go.Figure(data=data_traces, layout=layout)
    
    # Save HTML
    out_html = f"{output_name}.html"
    pyo.plot(fig, filename=out_html, auto_open=False)
    print(f"Saved {out_html}")

    # Save PDF
    if args.pdf:
        out_pdf = f"{output_name}.pdf"
        try:
            fig.write_image(out_pdf)
            print(f"Saved {out_pdf}")
        except Exception as e:
            print(f"Warning: Could not save PDF (requires 'pip install kaleido'): {e}")

def main():
    args = parse_args()
    
    # 1. Lattice & Special Points
    print(f"Reading lattice from {args.poscar}...")
    rec_cell, _ = get_reciprocal_lattice(args.poscar)
    special_points = get_special_points(args.poscar, rec_cell)

    # 2. Read Data
    print(f"Reading data from {args.input}...")
    if not Path(args.input).exists():
        print(f"Error: Input file '{args.input}' not found.")
        sys.exit(1)

    try:
        df = pd.read_csv(args.input, sep=r'\s+', comment='#', 
                         names=['qx', 'qy', 'qz', 'weight', 'band', 'freq', 'gamma'])
    except Exception as e:
        print(f"Error reading input file: {e}")
        sys.exit(1)

    # Filter by frequency
    df_filtered = df[df['freq'] > args.cut_f]
    if df_filtered.empty:
        print("Warning: No data points left after frequency filtering.")
        sys.exit(0)

    # 3. Process Data: Average & Minima
    print("Processing data (Average & Minima)...")
    grouped = df_filtered.groupby(['qx', 'qy', 'qz'])['gamma'].agg(['mean', 'min']).reset_index()
    
    q_frac = grouped[['qx', 'qy', 'qz']].values
    g_avg = grouped['mean'].values
    g_min = grouped['min'].values
    
    # Fix Position logic (Half BZ folding)
    if not args.no_fix:
        mask = q_frac[:, 2] <= -0.5
        q_frac[mask] *= -1

    q_cart = np.dot(q_frac, rec_cell)

    # Colorscale (Blue-White-Red custom for 0 crossing)
    if args.gmin < 0 < args.gmax:
        midpoint = (0 - args.gmin) / (args.gmax - args.gmin)
        colorscale = [[0.0, 'rgb(0, 0, 255)'], [midpoint, 'rgb(255, 255, 255)'], [1.0, 'rgb(255, 0, 0)']]
    else:
        colorscale = 'Plasma'

    # 4. Plot Average
    base_out = args.output
    if base_out.endswith(".html"): base_out = base_out[:-5]
    
    # Filter Average Data
    if args.hull_filter:
        print(f"Applying Convex Hull Filter (tol={args.hull_tol})...")
        q_cart_avg, g_avg_filt, q_frac_avg = apply_hull_filter(q_cart, g_avg, q_frac, special_points, args.hull_tol)
        q_cart_min, g_min_filt, q_frac_min = apply_hull_filter(q_cart, g_min, q_frac, special_points, args.hull_tol)
    else:
        q_cart_avg, g_avg_filt, q_frac_avg = q_cart, g_avg, q_frac
        q_cart_min, g_min_filt, q_frac_min = q_cart, g_min, q_frac

    plot_bz_with_data(q_cart_avg, q_frac_avg, g_avg_filt, rec_cell, special_points, "Average", f"{base_out}_averaged", args, colorscale)
    
    # 5. Plot Minima
    plot_bz_with_data(q_cart_min, q_frac_min, g_min_filt, rec_cell, special_points, "Minima", f"{base_out}_minima", args, colorscale)

    print("Done.")

if __name__ == "__main__":
    main()