import imgui
import numpy as np
import time
from typing import Dict, Tuple

try:
    import matplotlib.pyplot as plt
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False

def draw_rdf_plot(rdf_results: Dict[str, Tuple[np.ndarray, np.ndarray]], 
                  rdf_visibility: Dict[str, bool] = None,
                  normalize: bool = False):
    """Renders a unified scientific RDF plot with overlays, axis labels, and tick numbers. Supporting peak normalization."""
    imgui.set_next_window_bg_alpha(0.7)
    if not imgui.begin("Radial Distribution Explorer (g(r))"):
        imgui.end()
        return

    if not rdf_results:
        imgui.text("No data to plot. Use 'Compute RDF' in the analysis panel.")
        imgui.end()
        return

    if rdf_visibility is None: rdf_visibility = {k: True for k in rdf_results}

    # 1. Header & Actions
    imgui.text("Visible Species Pairs:")
    imgui.columns(3, border=False)
    for lbl in sorted(rdf_results.keys()):
        _, rdf_visibility[lbl] = imgui.checkbox(f"{lbl}", rdf_visibility.get(lbl, True))
        imgui.next_column()
    imgui.columns(1)
    
    imgui.spacing()
    if imgui.button(" Export Data (CSV) "):
        ts = int(time.time())
        active_lbls = [l for l in rdf_results if rdf_visibility.get(l, True)]
        for lbl in active_lbls:
            r, g = rdf_results[lbl]
            # Always export raw g(r) to CSV for integrity, but can be scaled later by user
            np.savetxt(f"rdf_{lbl}_{ts}.csv", np.column_stack((r, g)), header="r(A), g(r)", delimiter=",")
        print("RDF Data (Raw) exported to CSV.")
    
    imgui.same_line()
    if imgui.button(" Export Plot (PNG) "):
        if HAS_MATPLOTLIB:
            plt.figure(figsize=(10, 6))
            for lbl, (r, g) in rdf_results.items():
                if rdf_visibility.get(lbl, True):
                    g_plot = g / np.max(g) if normalize else g
                    plt.plot(r, g_plot, label=lbl + (" (Norm)" if normalize else ""), linewidth=2)
            plt.title("Radial Distribution Function" + (" (Normalized)" if normalize else ""))
            plt.xlabel("Distance r (Å)"); plt.ylabel("g(r) / max(g)" if normalize else "g(r)")
            plt.grid(True, alpha=0.3); plt.legend()
            plt.savefig(f"rdf_plot_{int(time.time())}.png", dpi=200)
            plt.close()
            print("Scientific plot saved.")

    imgui.separator(); imgui.spacing()

    # 2. Custom Drawing Engine
    p_orig = imgui.get_cursor_screen_pos()
    avail_w = imgui.get_content_region_available_width()
    plot_h = 320
    imgui.dummy(avail_w, plot_h)
    
    draw_list = imgui.get_window_draw_list()
    
    # Coordinates & Margins
    margin_l, margin_r, margin_t, margin_b = 70, 80, 20, 60
    p0 = (p_orig[0] + margin_l, p_orig[1] + margin_t)
    p1 = (p_orig[0] + avail_w - margin_r, p_orig[1] + plot_h - margin_b)
    cw, ch = (p1[0] - p0[0]), (p1[1] - p0[1])

    # Draw Background & Frame
    draw_list.add_rect_filled(p0[0], p0[1], p1[0], p1[1], 0x44000000)
    draw_list.add_rect(p0[0], p0[1], p1[0], p1[1], 0xFFAAAAAA)
    
    # Data scaling with normalization support
    max_g_all = 1e-9; max_r = 1e-9
    processed_results = {}
    for lbl, (r, g) in rdf_results.items():
        if not rdf_visibility.get(lbl, True) or r.size == 0: continue
        g_eff = g / np.max(g) if normalize else g
        processed_results[lbl] = (r, g_eff)
        max_g_all = max(max_g_all, np.max(g_eff))
        max_r = max(max_r, np.max(r))
    
    # 3. Y-Axis Labels
    draw_list.add_text(p0[0]-65, p0[1]+ch/2-10, 0xFF88CCFF, "Norm.g(r)" if normalize else "g(r)")
    draw_list.add_text(p0[0]-45, p0[1], 0xFFFFFFFF, f"{max_g_all:.1f}")
    draw_list.add_text(p0[0]-35, p1[1]-15, 0xFFFFFFFF, "0.0")

    # 4. X-Axis Ticks
    num_ticks = 6
    for t_idx in range(num_ticks):
        frac = t_idx / (num_ticks - 1)
        tx = p0[0] + frac * cw
        val = frac * max_r
        draw_list.add_line(tx, p1[1], tx, p1[1]+5, 0xFFAAAAAA)
        draw_list.add_text(tx-15, p1[1]+10, 0xFFFFFFFF, f"{val:.1f}")
    draw_list.add_text(p0[0] + cw/2 - 50, p1[1]+35, 0xFF88CCFF, "Distance r (Angstrom)")

    # 5. Drawing curves
    palette = [0xFF8888FF, 0xFF88FF88, 0xFFFFFF88, 0xFFFF88FF, 0xFF88FFFF, 0xFFCCCCCC]
    for i, (label, (r, g_eff)) in enumerate(processed_results.items()):
        color = palette[i % len(palette)]
        points = []
        for x_val, y_val in zip(r, g_eff):
            px = p0[0] + (x_val / max_r) * cw
            py = p1[1] - (y_val / max_g_all) * ch
            points.append((px, py))
        
        for p_idx in range(len(points) - 1):
             draw_list.add_line(points[p_idx][0], points[p_idx][1], points[p_idx+1][0], points[p_idx+1][1], color, 2.0)
             
        draw_list.add_text(p1[0] + 5, p0[1] + i*20, color, f"- {label}")

    imgui.end()

def draw_bad_plot(x: np.ndarray, y: np.ndarray):
    """Renders Bond Angle Distribution plot."""
    if not imgui.begin("Bond Angle Distribution"):
        imgui.end()
        return

    if x.size > 0:
        imgui.plot_histogram(
            "P(theta)",
            y.astype(np.float32),
            graph_size=(0, 150),
            scale_min=0.0
        )
        imgui.text("Angle (Degrees) 0 -> 180")
        
    imgui.end()
def draw_z_profile(x: np.ndarray, y: np.ndarray):
    """Renders Z-Density profile plot."""
    imgui.set_next_window_bg_alpha(0.7)
    if not imgui.begin("Z-Density Profile"):
        imgui.end()
        return

    if x.size > 0:
        imgui.plot_lines(
            "Density(z)",
            y.astype(np.float32),
            graph_size=(0, 150),
            scale_min=0.0
        )
        imgui.text(f"Z range: {x[0]:.2f} -> {x[-1]:.2f} A, bins: {len(x)}")
        
    imgui.end()

def draw_energy_plots(energies: np.ndarray, current_idx: int) -> int:
    """
    Renders high-quality, interactive energy analytics.
    Returns the new index if the user clicked on the trace, otherwise returns -1.
    """
    new_idx = -1
    imgui.set_next_window_bg_alpha(0.85)
    imgui.set_next_window_size(500, 480, imgui.FIRST_USE_EVER)
    if not imgui.begin("Energy Analytics Explorer", True):
        imgui.end()
        return new_idx

    if energies.size == 0:
        imgui.text_disabled("Waiting for structural data...")
        imgui.end()
        return new_idx

    # Guard: filter NaN/Inf - replace with interpolated or masked values
    valid_mask = np.isfinite(energies)
    if not np.any(valid_mask):
        imgui.text_colored("All energies are NaN/Inf — check data source.", 1.0, 0.4, 0.4)
        imgui.end()
        return new_idx
    
    # Replace invalid entries with nearest valid value for display purposes only
    energies_safe = energies.copy()
    if not np.all(valid_mask):
        valid_indices = np.where(valid_mask)[0]
        for i in np.where(~valid_mask)[0]:
            nearest = valid_indices[np.argmin(np.abs(valid_indices - i))]
            energies_safe[i] = energies[nearest]
        energies = energies_safe

    draw_list = imgui.get_window_draw_list()
    avail_w = imgui.get_content_region_available_width()
    p_orig = imgui.get_cursor_screen_pos()
    
    # --- Part 1: Energy Trace ---
    imgui.text_colored("Energy Landscape (Trajectory)", 0.6, 0.8, 1.0)
    plot_h = 200
    margin_l, margin_r, margin_t, margin_b = 65, 25, 10, 35
    p0 = (p_orig[0] + margin_l, p_orig[1] + margin_t)
    p1 = (p_orig[0] + avail_w - margin_r, p_orig[1] + plot_h - margin_b)
    cw, ch = (p1[0] - p0[0]), (p1[1] - p0[1])

    # Invisible button for overall interaction
    imgui.invisible_button("##trace_zone", avail_w, plot_h)
    is_hovered = imgui.is_item_hovered()
    
    # Draw Background
    draw_list.add_rect_filled(p0[0], p0[1], p1[0], p1[1], 0x55111111, rounding=4)
    draw_list.add_rect(p0[0], p0[1], p1[0], p1[1], 0xAA666666, rounding=4)

    e_min, e_max = np.min(energies), np.max(energies)
    e_range = e_max - e_min if e_max > e_min else 1.0
    e_mean = np.mean(energies)

    # Draw Trace Line & Area Fill
    num_pts = min(len(energies), int(cw))
    if num_pts > 1:
        step = len(energies) / num_pts
        points = []
        for i in range(num_pts):
            idx = int(i * step)
            val = energies[idx]
            px = p0[0] + (i / (num_pts - 1)) * cw
            py = p1[1] - ((val - e_min) / e_range) * ch
            points.append((px, py))
        
        # Area Fill
        for i in range(len(points) - 1):
             draw_list.add_rect_filled(
                 points[i][0], points[i][1], points[i+1][0], p1[1],
                 0x3388CCFF
             )
        # Main Line
        for i in range(len(points) - 1):
            draw_list.add_line(points[i][0], points[i][1], points[i+1][0], points[i+1][1], 0xFFAABBFF, 2.0)

    # Draw Mean Line (Dashed-style)
    mean_y = p1[1] - ((e_mean - e_min) / e_range) * ch
    draw_list.add_line(p0[0], mean_y, p1[0], mean_y, 0x88FFFFFF, 1.0)
    draw_list.add_text(p1[0] + 5, mean_y - 7, 0x88FFFFFF, "Mean")

    # Current Frame Indicator
    curr_norm = current_idx / (len(energies)-1) if len(energies) > 1 else 0.5
    cx = p0[0] + curr_norm * cw
    cy = p1[1] - ((energies[current_idx] - e_min) / e_range) * ch
    draw_list.add_line(cx, p0[1], cx, p1[1], 0x44FF00FF, 1.0) # Vertical guide
    draw_list.add_circle_filled(cx, cy, 6.0, 0xFFFF00FF)
    draw_list.add_circle(cx, cy, 8.0, 0xFFFFFFFF, num_segments=16)

    # Hover Interactivity (Tooltip + Vertical Line)
    if is_hovered:
        m_pos = imgui.get_mouse_pos()
        if p0[0] <= m_pos[0] <= p1[0]:
            h_idx = int(((m_pos[0] - p0[0]) / cw) * (len(energies)-1))
            h_idx = np.clip(h_idx, 0, len(energies)-1)
            hx = p0[0] + (h_idx / (len(energies)-1)) * cw
            hy = p1[1] - ((energies[h_idx] - e_min) / e_range) * ch
            
            draw_list.add_line(hx, p0[1], hx, p1[1], 0xAAFFFFFF, 1.0)
            draw_list.add_circle_filled(hx, hy, 4.0, 0xFFFFFFFF)
            
            imgui.set_tooltip(f"Frame: {h_idx}\nEnergy: {energies[h_idx]:.5f} eV")
            
            if imgui.is_mouse_clicked(0):
                new_idx = h_idx

    # Labels
    draw_list.add_text(p0[0] - 60, p0[1], 0xFFFFFFFF, f"{e_max:.2f}")
    draw_list.add_text(p0[0] - 60, p1[1] - 15, 0xFFFFFFFF, f"{e_min:.2f}")
    draw_list.add_text(p0[0] + cw/2 - 40, p1[1] + 10, 0xFF88CCFF, "Trajectory Path")

    # --- Part 2: Distribution Histogram ---
    imgui.set_cursor_screen_pos((p_orig[0], p_orig[1] + plot_h + 20))
    imgui.text_colored("Energy Distribution (Density)", 0.6, 1.0, 0.8)
    
    hist_h = 160
    p_hist_orig = imgui.get_cursor_screen_pos()
    imgui.dummy(avail_w, hist_h)
    
    h_p0 = (p_hist_orig[0] + margin_l, p_hist_orig[1] + 10)
    h_p1 = (p_hist_orig[0] + avail_w - margin_r, p_hist_orig[1] + hist_h - 30)
    hcw, hch = (h_p1[0] - h_p0[0]), (h_p1[1] - h_p0[1])

    draw_list.add_rect_filled(h_p0[0], h_p0[1], h_p1[0], h_p1[1], 0x55111111, rounding=4)
    
    num_bins = min(40, len(energies)//4 + 5)
    counts, bins = np.histogram(energies, bins=num_bins)
    c_max = np.max(counts) if counts.size > 0 else 1
    
    # Find bin for current energy
    curr_e = energies[current_idx]
    curr_bin = np.digitize(curr_e, bins) - 1

    if counts.size > 0:
        bin_w = hcw / counts.size
        for i, count in enumerate(counts):
            bx0 = h_p0[0] + i * bin_w
            bx1 = bx0 + bin_w - 1
            by0 = h_p1[1] - (count / c_max) * hch
            by1 = h_p1[1]
            
            # Use gradient/different color for the "current" bin
            col = 0xFF88FF88 if i != curr_bin else 0xFFFF8888
            draw_list.add_rect_filled(bx0, by0, bx1, by1, col)
            if count > 0:
                draw_list.add_rect(bx0, by0, bx1, by1, 0x44FFFFFF)

    # Summary Stats
    imgui.set_cursor_screen_pos((p_hist_orig[0] + 10, p_hist_orig[1] + hist_h - 10))
    imgui.text_disabled(f"Avg: {e_mean:.5f} | Std: {np.std(energies):.5f} | N: {len(energies)}")

    imgui.end()
    return new_idx


# ────────────────────────────────────────────────
# Tight-Binding DOS / PDOS Plot
# ────────────────────────────────────────────────
_SPECIES_COLORS = [
    0xFF4488FF, 0xFFFF6644, 0xFF44CC77, 0xFFCC44EE,
    0xFFFFCC33, 0xFF33CCFF, 0xFFFF4477, 0xFF77FF44,
]

def draw_dos_plot(dos_x: np.ndarray, dos_y: np.ndarray, fermi: float,
                  pdos_y: dict = None, homo: float = 0.0, lumo: float = 0.0,
                  gap: float = 0.0):
    """
    Renders an interactive DOS / PDOS plot with Fermi level marker.
    dos_x, dos_y: total DOS arrays
    pdos_y: dict{species: ndarray} of projected DOS
    fermi: Fermi level energy (after reference shift)
    """
    imgui.set_next_window_bg_alpha(0.85)
    imgui.set_next_window_size(440, 380, imgui.FIRST_USE_EVER)
    if not imgui.begin("Electron DOS (Tight-Binding)"):
        imgui.end()
        return

    if dos_x is None or len(dos_x) == 0:
        imgui.text_disabled("Run TB calculation to see DOS.")
        imgui.end()
        return

    draw_list = imgui.get_window_draw_list()
    avail_w = imgui.get_content_region_available_width()
    plot_h = 220
    p0 = imgui.get_cursor_screen_pos()
    p1 = (p0[0] + avail_w, p0[1] + plot_h)

    # Background
    draw_list.add_rect_filled(p0[0], p0[1], p1[0], p1[1], 0xFF1A1A2A)
    draw_list.add_rect(p0[0], p0[1], p1[0], p1[1], 0xFF555577)

    e_min, e_max = float(dos_x[0]), float(dos_x[-1])
    d_max = float(dos_y.max()) * 1.05 if dos_y.max() > 0 else 1.0
    w = p1[0] - p0[0]
    h = p1[1] - p0[1]

    def ex(e): return p0[0] + (e - e_min) / (e_max - e_min) * w
    def ey(d): return p1[1] - (d / d_max) * h

    # Total DOS fill
    pts = [(ex(dos_x[0]), p1[1])]
    for i in range(len(dos_x)):
        pts.append((ex(dos_x[i]), ey(dos_y[i])))
    pts.append((ex(dos_x[-1]), p1[1]))
    for k in range(len(pts) - 1):
        draw_list.add_line(pts[k][0], pts[k][1], pts[k+1][0], pts[k+1][1], 0x6644AAFF, 1.5)

    # PDOS curves
    if pdos_y:
        for ci, (sp, py) in enumerate(pdos_y.items()):
            col = _SPECIES_COLORS[ci % len(_SPECIES_COLORS)]
            for i in range(len(dos_x) - 1):
                draw_list.add_line(
                    ex(dos_x[i]),   ey(float(py[i])),
                    ex(dos_x[i+1]), ey(float(py[i+1])),
                    col, 1.5
                )

    # Fermi level (vertical dashed line)
    fx = ex(fermi)
    if p0[0] <= fx <= p1[0]:
        for yy in range(int(p0[1]), int(p1[1]), 6):
            draw_list.add_line(fx, yy, fx, min(yy+3, p1[1]), 0xFFFFDD44, 1.5)
        draw_list.add_text(fx + 3, p0[1] + 4, 0xFFFFDD44, "Ef")

    # Axis ticks (energy)
    n_ticks = 5
    for k in range(n_ticks + 1):
        e_tick = e_min + k * (e_max - e_min) / n_ticks
        xt = ex(e_tick)
        draw_list.add_line(xt, p1[1] - 4, xt, p1[1], 0x88FFFFFF)
        draw_list.add_text(xt - 12, p1[1] + 2, 0xAAFFFFFF, f"{e_tick:.1f}")

    # Advance cursor past plot
    imgui.set_cursor_screen_pos((p0[0], p1[1] + 20))
    imgui.dummy(avail_w, 0.0)

    # Metrics
    imgui.text_colored(f"HOMO: {homo:+.3f} eV", 0.4, 1.0, 0.6)
    imgui.same_line()
    imgui.text_colored(f"  LUMO: {lumo:+.3f} eV", 1.0, 0.6, 0.4)
    imgui.same_line()
    imgui.text_colored(f"  Gap: {gap:.3f} eV", 1.0, 1.0, 0.4)

    # Legend for PDOS
    if pdos_y:
        imgui.separator()
        imgui.text("PDOS:")
        for ci, sp in enumerate(pdos_y):
            r = (((_SPECIES_COLORS[ci % len(_SPECIES_COLORS)] >> 0)  & 0xFF)) / 255.0
            g = (((_SPECIES_COLORS[ci % len(_SPECIES_COLORS)] >> 8)  & 0xFF)) / 255.0
            b = (((_SPECIES_COLORS[ci % len(_SPECIES_COLORS)] >> 16) & 0xFF)) / 255.0
            imgui.text_colored(f"  ■ {sp}", r, g, b)
            imgui.same_line()

    imgui.end()


# ────────────────────────────────────────────────
# Eigenvalue Ladder Plot
# ────────────────────────────────────────────────
def draw_eigenvalue_ladder(eigenvalues: np.ndarray, n_occ: int,
                           ipr: np.ndarray = None, fermi: float = 0.0):
    """
    Renders an energy level diagram. Occupied states (green), unoccupied (red).
    Colors by IPR if provided (localized=orange, delocalized=blue).
    """
    imgui.set_next_window_bg_alpha(0.85)
    imgui.set_next_window_size(200, 360, imgui.FIRST_USE_EVER)
    if not imgui.begin("Energy Levels (TB)"):
        imgui.end()
        return

    if eigenvalues is None or len(eigenvalues) == 0:
        imgui.text_disabled("No eigenvalues.")
        imgui.end()
        return

    draw_list = imgui.get_window_draw_list()
    avail_w = imgui.get_content_region_available_width()
    plot_h = 300
    p0 = imgui.get_cursor_screen_pos()

    draw_list.add_rect_filled(p0[0], p0[1], p0[0] + avail_w, p0[1] + plot_h, 0xFF1A1A2A)

    # Show max 80 levels centered around HOMO/LUMO
    center = n_occ
    half = 40
    lo = max(0, center - half)
    hi = min(len(eigenvalues), center + half)
    ev_show = eigenvalues[lo:hi]

    e_min = float(ev_show.min()) - 0.5
    e_max = float(ev_show.max()) + 0.5
    h = float(plot_h)

    def ey(e): return p0[1] + (1 - (e - e_min) / (e_max - e_min)) * h

    for idx, e in enumerate(ev_show):
        global_idx = lo + idx
        y = ey(float(e))
        occ = global_idx < n_occ

        # Color: occupied=green, unoccupied=dim red; modulated by IPR
        if ipr is not None and len(ipr) > global_idx:
            loc = min(1.0, float(ipr[global_idx]) * len(eigenvalues))
            col = (
                0xFF0088FF if not occ else
                (0xFF2266FF if loc < 0.3 else 0xFFFF8822 if loc > 0.7 else 0xFF44CC44)
            )
        else:
            col = 0xFF44CC44 if occ else 0xFF446688

        draw_list.add_line(p0[0] + 20, y, p0[0] + avail_w - 20, y, col, 1.5)

    # Fermi level
    fy = ey(fermi)
    if p0[1] <= fy <= p0[1] + plot_h:
        draw_list.add_line(p0[0], fy, p0[0] + avail_w, fy, 0xFFFFDD44, 2.0)

    imgui.set_cursor_screen_pos((p0[0], p0[1] + plot_h + 4))
    imgui.text_disabled(f"{lo}–{hi-1} of {len(eigenvalues)} states")
    imgui.end()
