#!/usr/bin/env python3
import argparse
import os
import sys
import yaml
import numpy as np
import matplotlib.pyplot as plt
from ase.io import read
from pathlib import Path

# Constants
KB_J = 1.380649e-23     # Boltzmann constant in J/K
E_CHARGE = 1.60217663e-19 # Elementary charge in C
ANGSTROM = 1e-10        # m
FEMTOSECOND = 1e-15     # s

def load_default_charges():
    """Load default oxidation states from pydefect database if available."""
    # 1. Try to find via pydefect package
    try:
        import pydefect
        pydefect_path = Path(pydefect.__file__).parent
        yaml_path = pydefect_path / "database" / "oxidation_state.yaml"
        if yaml_path.exists():
            with open(yaml_path, "r") as f:
                return yaml.safe_load(f)
    except (ImportError, Exception):
        pass

    # 2. Try relative to script (original logic)
    script_path = Path(__file__).resolve()
    project_root = script_path.parent.parent
    yaml_path = project_root / "pydefect" / "database" / "oxidation_state.yaml"
    
    if yaml_path.exists():
        try:
            with open(yaml_path, "r") as f:
                return yaml.safe_load(f)
        except Exception as e:
            print(f"Warning: Could not load oxidation states from {yaml_path}: {e}")
    return {}

def parse_user_charges(user_str, defaults):
    """Update defaults with user provided string 'El:Q,El2:Q2'."""
    charges = defaults.copy()
    if not user_str:
        return charges
    
    parts = user_str.split(',')
    for p in parts:
        p = p.strip()
        if ':' in p:
            el, q = p.split(':')
            try:
                charges[el.strip()] = float(q)
            except ValueError:
                print(f"Warning: Could not parse charge for '{p}'")
    return charges

def detect_interval(input_path):
    """Try to detect interval from XDATCAR configuration numbers."""
    if 'XDATCAR' not in str(input_path):
        return None
    try:
        configs = []
        with open(input_path, 'r') as f:
            for line in f:
                if "Direct configuration=" in line:
                    val = line.split('=')[1].strip().split()[0]
                    configs.append(int(val))
                if len(configs) >= 2:
                    break
        if len(configs) >= 2:
            return configs[1] - configs[0]
    except Exception as e:
        pass
    return None

def calculate_conductivity(args):
    input_path = Path(args.input)
    if not input_path.exists():
        print(f"Error: Input file {args.input} not found.")
        return

    # Auto-detect interval if not provided (or set to 1)
    if args.interval == 1:
        detected = detect_interval(input_path)
        if detected:
            args.interval = detected
            print(f"Auto-detected interval: {args.interval}")

    # Setup output path (in the same directory as input)
    out_prefix = input_path.parent / args.out
    log_file = f"{out_prefix}.log"
    
    def log_print(msg):
        print(msg)
        with open(log_file, "a") as f:
            f.write(msg + "\n")
            
    # Clear log file if exists
    with open(log_file, "w") as f:
        f.write("--- MD Conductivity Analysis ---\n")
        f.write(f"Input: {args.input}\n")
        f.write(f"Temp: {args.temp} K\n")
        f.write(f"Time step: {args.dt} fs\n")
        f.write(f"Interval: {args.interval}\n")
        if args.charge_msd:
            f.write("Mode: Charge MSD (correlated) enabled\n")
        f.write("-" * 40 + "\n")

    log_print(f"Reading trajectory: {args.input}")
    
    try:
        # Auto-detect format: use 'vasp-xdatcar' only if 'XDATCAR' is in the name
        fmt = 'vasp-xdatcar' if 'XDATCAR' in str(args.input) else None
        from ase.io import iread, read
        import io
        traj = []
        
        try:
            for atoms in iread(str(args.input), format=fmt):
                traj.append(atoms)
        except Exception as e:
            if fmt == 'vasp-xdatcar':
                log_print(f"Standard XDATCAR reading failed: {e}")
                log_print("Attempting robust reading mode for concatenated XDATCAR...")
                with open(str(args.input), 'r') as f:
                    content = f.read()
                
                if not content.strip():
                    raise ValueError("File is empty")
                    
                first_line = content.splitlines()[0]
                pieces = content.split(first_line + '\n')
                traj = []
                for p in pieces:
                    if not p.strip():
                        continue
                    try:
                        frames = read(io.StringIO(first_line + '\n' + p), index=':', format='vasp-xdatcar')
                        traj.extend(frames)
                    except:
                        continue
                if not traj:
                    raise e
                log_print(f"Successfully read {len(traj)} frames in robust mode.")
            else:
                raise e
    except Exception as e:
        log_print(f"Error reading trajectory: {e}")
        return

    if not traj:
        log_print("Error: Trajectory is empty.")
        return

    n_steps = len(traj)
    n_atoms = len(traj[0])
    species = sorted(list(set(traj[0].get_chemical_symbols())))
    
    vol_A3 = traj[0].get_volume()
    vol_m3 = vol_A3 * (ANGSTROM ** 3)
    
    log_print(f"Steps: {n_steps}, Atoms: {n_atoms}, Volume: {vol_A3:.2f} A^3")
    log_print(f"Species found: {species}")
    
    # Charges
    defaults = load_default_charges()
    charges = parse_user_charges(args.charges, defaults)
    
    log_print("Using oxidation states:")
    for sp in species:
        q = charges.get(sp, None)
        if q is None:
            q = 0.0
            log_print(f"  {sp}: {q} (Warning: not found in defaults, using 0.0)")
        else:
            log_print(f"  {sp}: {q}")
        charges[sp] = q
        
    # Unwrapping
    log_print("Unwrapping coordinates...")
    pos_list = [atoms.get_scaled_positions() for atoms in traj]
    scaled_pos = np.stack(pos_list)
    d_scaled = scaled_pos[1:] - scaled_pos[:-1]
    # Apply Minimum Image Convention
    d_scaled -= np.round(d_scaled)
    
    cum_disp_scaled = np.zeros_like(scaled_pos)
    cum_disp_scaled[1:] = np.cumsum(d_scaled, axis=0)
    
    # Handle Cell (Variable or Constant)
    cell_0 = traj[0].get_cell()
    cell_last = traj[-1].get_cell()
    if np.allclose(cell_0, cell_last, atol=1e-3):
        cell = cell_0
        unwrapped_real = np.dot(cum_disp_scaled + scaled_pos[0], cell)
    else:
        log_print("Variable cell detected. Transforming per step (approximation)...")
        unwrapped_real = np.zeros_like(scaled_pos)
        # Project accumulated scaled positions onto the instantaneous cell
        for t in range(n_steps):
            unwrapped_real[t] = (cum_disp_scaled[t] + scaled_pos[0]) @ traj[t].get_cell()

    # Effective time step
    dt_effective = args.dt * args.interval
    time_ps = np.arange(n_steps) * dt_effective * 1e-3

    # MSD & Sigma
    results = {}
    
    # --------------------------------------------------
    # Charge MSD preparation (Einstein form)
    # --------------------------------------------------
    if args.charge_msd:
        log_print("\nComputing charge MSD (Einstein form)...")

        symbols = traj[0].get_chemical_symbols()
        charges_SI = np.zeros(n_atoms)

        for i, sp in enumerate(symbols):
            charges_SI[i] = charges.get(sp, 0.0) * E_CHARGE

        # q_i * (r_i(t) - r_i(0))
        disp_charge = np.zeros((n_steps, 3))
        for t in range(n_steps):
            dr = unwrapped_real[t] - unwrapped_real[0]
            disp_charge[t] = np.sum(dr * charges_SI[:, None], axis=0)

        # | Σ q_i (r_i(t) - r_i(0)) |^2
        charge_msd = np.sum(disp_charge**2, axis=1)

    plt.figure(figsize=(8, 6))
    
    for sp in species:
        indices = [i for i, s in enumerate(traj[0].get_chemical_symbols()) if s == sp]
        if not indices:
            continue
            
        n_ions = len(indices)
        pos_sp = unwrapped_real[:, indices, :]
        
        # MSD(t) = < |r(t) - r(0)|^2 >
        disp_from_0 = pos_sp - pos_sp[0]
        sq_dist = np.sum(disp_from_0**2, axis=2)
        msd = np.mean(sq_dist, axis=1)
        
        # Fit (10% to 90% range to avoid initial ballistic and poor statistics at end)
        start_idx = int(n_steps * 0.1)
        end_idx = int(n_steps * 0.9)
        if start_idx >= end_idx:
            start_idx, end_idx = 0, n_steps
            
        fit_time = time_ps[start_idx:end_idx]
        fit_msd = msd[start_idx:end_idx]
        
        slope = 0
        intercept = 0
        if len(fit_time) > 1:
            slope, intercept = np.polyfit(fit_time, fit_msd, 1)
            
        # D in cm^2/s
        # MSD = 6Dt => D = slope / 6
        D_A2_ps = slope / 6.0
        D_SI = D_A2_ps * 1e-8 # m^2/s (1e-20 / 1e-12)
        D_cm2_s = D_SI * 1e4
        
        # Sigma in S/m
        # sigma = (N/V) * (q^2 * D) / (kT)
        q = charges.get(sp, 0.0) * E_CHARGE
        n_density = n_ions / vol_m3
        
        if q == 0:
            sigma = 0.0
        else:
            sigma = (n_density * (q**2) * D_SI) / (KB_J * args.temp)
            
        results[sp] = {'D': D_cm2_s, 'sigma': sigma, 'msd': msd}
        
        log_print(f"\n--- Species: {sp} ---")
        log_print(f"  Count: {n_ions}")
        log_print(f"  Diff. Coeff (D): {D_cm2_s:.3e} cm^2/s")
        log_print(f"  Conductivity (sigma_NE): {sigma:.3e} S/m ({sigma*10:.2f} mS/cm)")
        
        # Plotting
        plt.plot(time_ps, msd, label=f"{sp} ($D$={D_cm2_s:.1e})")
        if len(fit_time) > 1:
            plt.plot(fit_time, slope * fit_time + intercept, '--', alpha=0.5, color='gray')

    # --------------------------------------------------
    # Ionic conductivity from charge MSD (correlated)
    # --------------------------------------------------
    sigma_charge = None

    if args.charge_msd:
        start_idx = int(n_steps * 0.1)
        end_idx = int(n_steps * 0.9)
        if start_idx >= end_idx:
            start_idx, end_idx = 0, n_steps

        fit_time = time_ps[start_idx:end_idx]
        fit_charge_msd = charge_msd[start_idx:end_idx]

        if len(fit_time) > 1:
            slope_charge, _ = np.polyfit(fit_time, fit_charge_msd, 1)
        else:
            slope_charge = 0.0

        # (C·Å)^2 → (C·m)^2
        slope_charge_SI = slope_charge * (ANGSTROM ** 2)

        sigma_charge = slope_charge_SI / (6 * vol_m3 * KB_J * args.temp)

        log_print("\n--- Charge MSD conductivity (correlated) ---")
        log_print(f"  sigma_charge = {sigma_charge:.3e} S/m ({sigma_charge*10:.2f} mS/cm)")

    total_sigma = sum(r['sigma'] for r in results.values())

    log_print("\n===== Conductivity Summary =====")
    log_print(f"Total sigma_NE: {total_sigma:.3e} S/m ({total_sigma*10:.2f} mS/cm)")

    if sigma_charge is not None:
        haven = sigma_charge / total_sigma if total_sigma > 0 else 0.0
        log_print(f"Sigma from charge MSD: {sigma_charge:.3e} S/m ({sigma_charge*10:.2f} mS/cm)")
        log_print(f"Haven ratio (sigma_charge / sigma_NE): {haven:.3f}")
    
    plt.xlabel("Time (ps)")
    plt.ylabel(r"MSD ($\AA^2$)")
    plt.title(rf"MSD @ {args.temp}K (Total $\sigma$ = {total_sigma*10:.2f} mS/cm)")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    
    plot_file = f"{out_prefix}.pdf"
    plt.savefig(plot_file)
    log_print(f"Plot saved to: {plot_file}")

    if args.charge_msd:
        plt.figure(figsize=(8, 6))
        plt.plot(time_ps, charge_msd)
        plt.xlabel("Time (ps)")
        plt.ylabel(r"Charge MSD $(C\cdot \AA)^2$")
        plt.title("Charge MSD (Einstein conductivity)")
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(f"{out_prefix}_charge_msd.pdf")
        log_print(f"Charge MSD plot saved to: {out_prefix}_charge_msd.pdf")

    log_print(f"Log saved to: {log_file}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Calculate MSD and Ionic Conductivity from trajectory (XDATCAR or md.traj).",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("input", nargs="?", default="XDATCAR", help="Path to trajectory file (XDATCAR or md.traj)")
    parser.add_argument("--dt", type=float, default=2.0, help="MD Time step in fs")
    parser.add_argument("--temp", type=float, default=600.0, help="Temperature in K")
    parser.add_argument("--interval", type=int, default=1, help="Sampling interval (stride) of XDATCAR vs MD steps")
    parser.add_argument("--out", default="md_results", help="Output prefix for log and plot files")
    parser.add_argument("--charges", type=str, default="", help="Override oxidation states, e.g., 'Li:1,P:5,S:-2'")
    parser.add_argument(
        "--charge_msd",
        action="store_true",
        help="Also compute ionic conductivity from charge MSD (Einstein form)."
    )
    
    args = parser.parse_args()
    
    calculate_conductivity(args)