import xml.etree.ElementTree as ET
import pandas as pd
import matplotlib.pyplot as plt
# Create the figure with a specific DPI
fig = plt.figure(dpi=1200)
import os
import numpy as np

def parse_xcd_2d_chart_multi_series(file_path):
    """Parses 2D chart data from a Materials Studio .xcd file with multiple series."""
    tree = ET.parse(file_path)
    root = tree.getroot()

    data = {}
    for series in root.findall('.//SERIES_2D'):
        name = series.attrib.get("Name", "Unnamed Series")
        points = []
        for pt in series.findall('POINT_2D'):
            xy = pt.attrib.get("XY")
            if xy:
                try:
                    k, energy = map(float, xy.split(','))
                    points.append((k, energy))
                except ValueError:
                    continue
        if points:
            df = pd.DataFrame(points, columns=["k-point", "Energy (eV)"])
            data[name] = df
    return data

def analyze_band_gap(band_series):
    """Find VBM, CBM, band gap, and determine direct vs indirect."""
    vbm = -np.inf
    cbm = np.inf
    vbm_k = cbm_k = None
    fermi = 0.0

    for name, df in band_series.items():
        val = df[df["Energy (eV)"] <= fermi]
        cond = df[df["Energy (eV)"] > fermi]

        if not val.empty:
            e_vbm = val["Energy (eV)"].max()
            if e_vbm > vbm:
                vbm = e_vbm
                vbm_k = val.loc[val["Energy (eV)"].idxmax(), "k-point"]

        if not cond.empty:
            e_cbm = cond["Energy (eV)"].min()
            if e_cbm < cbm:
                cbm = e_cbm
                cbm_k = cond.loc[cond["Energy (eV)"].idxmin(), "k-point"]

    gap = cbm - vbm
    direct = abs(vbm_k - cbm_k) < 1e-6
    return {
        "vbm_energy": vbm,
        "vbm_k_point": vbm_k,
        "cbm_energy": cbm,
        "cbm_k_point": cbm_k,
        "band_gap": gap,
        "band_type": "Direct" if direct else "Indirect"
    }

def shift_and_clip_pdos(pdos_series, vbm, cbm):
    """Shift PDOS Y by VBM→0 eV, and zero DOS intensity inside the gap [vbm, cbm]."""
    out = {}
    for name, df in pdos_series.items():
        df2 = df.copy()
        df2["Energy (eV)"] = df2["Energy (eV)"] - vbm
        mask = (df2["Energy (eV)"] > 0) & (df2["Energy (eV)"] < (cbm - vbm))
        df2.loc[mask, "k-point"] = 0.0
        out[name] = df2
    return out

def has_valid_pdos_data(df):
    """Check if PDOS data frame has meaningful DOS intensity."""
    if df.empty:
        return False
    dos_values = df["k-point"]
    if np.allclose(dos_values, 0):
        return False
    if np.allclose(df["Energy (eV)"], df["Energy (eV)"].iloc[0]):
        return False
    return True

def is_f_orbital(series_name):
    """Check if the series name indicates f-orbital data.
    This function is robust in detecting any reference to f-orbitals,
    even if they are unexpected in s-p-d datasets."""
    
    # Common patterns for f-orbital labels
    f_orbital_indicators = [
        'f-orbital', 'f orbital', 'f_orbital', 'orbital-f', 
        'orbital_f', 'orbital f', ' f ', '-f-', '_f_', ' F ', 
        'F-orbital', 'F orbital', 'F_orbital', 'orbital f',
        'f(', 'f)', 'f-', '-f', 'f_', '_f'
    ]
    
    name_lower = series_name.lower()
    
    # Standard check for f-orbital indicators
    if any(indicator.lower() in name_lower for indicator in f_orbital_indicators):
        return True
    
    # Additional check for orbital naming schemes like "Au 4f", "Rb f", etc.
    orbital_patterns = ['f', 'F']
    for pattern in orbital_patterns:
        if pattern in series_name:
            # Avoid matching cases where 'f' is just part of another word
            word_boundary = [' '+pattern+' ', ' '+pattern, pattern+' ', '-'+pattern, pattern+'-']
            if any(wb in series_name for wb in word_boundary):
                return True
    
    return False

def plot_series_horizontal_subplots(data):
    """Plot band structure & PDOS side by side, aligned and with identical gap."""
    klabels = {0.0: "Γ", 0.0204: "A", 0.299: "H", 0.32: "K",
               0.599: "Γ", 0.84: "M", 0.861: "L", 1.0: "H"}

    bands = {n: df for n, df in data.items() if "band" in n.lower()}
    pdoss = {n: df for n, df in data.items() if "band" not in n.lower()}

    if not bands:
        print("No band-series found.")
        return

    # Analyze gap
    gap = analyze_band_gap(bands)
    print("VBM: {:.4f} eV at k={:.4f}".format(gap["vbm_energy"], gap["vbm_k_point"]))
    print("CBM: {:.4f} eV at k={:.4f}".format(gap["cbm_energy"], gap["cbm_k_point"]))
    print("Eg:  {:.4f} eV ({})".format(gap["band_gap"], gap["band_type"]))

    # Shift & clip PDOS
    pdoss_shifted = shift_and_clip_pdos(pdoss, gap["vbm_energy"], gap["cbm_energy"])

    y_min, y_max = -10, 10

    fig, (axb, axp) = plt.subplots(1, 2, figsize=(6, 3), dpi=300, sharey=True)

    # ----- BAND STRUCTURE -----
    for name, df in bands.items():
        y = df["Energy (eV)"] - gap["vbm_energy"]
        axb.plot(df["k-point"], y, linewidth=1.2)

    # Mark VBM and CBM with draggable legends
    vbm_marker, = axb.plot(gap["vbm_k_point"], 0, 'go', ms=3, label="VBM")
    cbm_marker, = axb.plot(gap["cbm_k_point"], gap["band_gap"], 'bo', ms=3, label="CBM")
    axb.axhline(0, color='red', ls=':', lw=0.5)

    # Draggable VBM/CBM legend
    lg1 = axb.legend(handles=[vbm_marker, cbm_marker], fontsize=5, frameon=False, loc='best')
    lg1.set_draggable(True)
    axb.add_artist(lg1)

    # Add Eg annotation as a draggable legend entry
    eg_text = f"Eg = {gap['band_gap']:.2f} eV ({gap['band_type']})"
    dummy_line, = axb.plot([], [], ' ', label=eg_text)  # empty handle just for legend
    lg2 = axb.legend(handles=[dummy_line], fontsize=5, frameon=False, loc='center')
    lg2.set_draggable(True)

    axb.set_title("Band Structure", fontsize=8)
    axb.set_ylabel("Energy (eV)", fontsize=6)
    axb.set_xlim(0, 1)
    axb.set_ylim(y_min, y_max)
    axb.set_xticks(list(klabels.keys()))
    axb.set_xticklabels(list(klabels.values()), fontsize=5)
    axb.tick_params(labelsize=5)

    # ----- PARTIAL DOS -----
    lines, labels = [], []
    
    # For debugging purposes
    debug_mode = False  # Set to True to see which orbital series are filtered out
    if debug_mode:
        print("PDOS series and their filtering status:")
    
    for name, df in pdoss_shifted.items():
        # Skip f-orbitals for the legend
        if is_f_orbital(name):
            if debug_mode:
                print(f"Filtered out: {name}")
                
            # Still plot the data, but don't add to legend
            axp.plot(df["k-point"], df["Energy (eV)"], lw=1.2)
            continue
        
        if debug_mode:
            print(f"Included: {name}")
            
        if has_valid_pdos_data(df):
            ln, = axp.plot(df["k-point"], df["Energy (eV)"], lw=1.2, label=name)
            lines.append(ln)
            labels.append(name)

    axp.set_xlim(0, 20)
    axp.set_ylim(y_min, y_max)
    axp.axhline(0, color='red', ls=':', lw=0.5)
    axp.set_xlabel("Density of States (electron/eV)", fontsize=6)
    axp.set_title("Partial DOS", fontsize=8)
    axp.tick_params(labelsize=5)

    if lines:
        lg3 = axp.legend(lines, labels, fontsize=5, frameon=False, loc='best')
        lg3.set_draggable(True)

    plt.tight_layout(rect=[0, 0, 1, 0.95])
    plt.show()

def convert_xcd_to_csv_and_plot(xcd_path):
    if not os.path.exists(xcd_path):
        raise FileNotFoundError(f"File not found: {xcd_path}")
    data = parse_xcd_2d_chart_multi_series(xcd_path)
    
    # Debug: print all series names to help identify the problematic f-orbital series
    #print("All series in the file:")
    #for name in data.keys():
    #print(f"  - {name}")
    
    plot_series_horizontal_subplots(data)

if __name__ == "__main__":
    xcd_file = r"C:\Users\User\Documents\Materials Studio Projects\3D-RbAuSe_Files\Documents\bulk-RbAuSe CASTEP GeomOpt\bulk-RbAuSe Band Structure (2).xcd"
    convert_xcd_to_csv_and_plot(xcd_file)
ı have script likr this generating two plots. I would like to have different colors for valence and conduction bands in the first plot by keeping everthing else same (unchanged)