"""
in-situ TEM Toolbox - DENS holder data process

Assembles of functions related to process raw data by DENS holders

Example:
    
    from insitu_TEM import insitu_DENS as DENS
    DENS.XXX

Created on Tue May 13 2025
@author: Meng Li mli4@bnl.gov

"""

import pandas as pd

def read_pstrace_csv(filepath):
    # Try several encodings until one works
    for encoding in ['utf-16', 'utf-8', 'ISO-8859-1']:
        try:
            with open(filepath, 'r', encoding=encoding) as f:
                lines = [line.strip() for line in f if line.strip()]
            break
        except UnicodeDecodeError:
            continue
    else:
        raise UnicodeDecodeError("Failed to decode file using common encodings.")

    # Detect measurement type (e.g., CV I vs E Scan 1)
    measurement_type = None
    for line in lines:
        if "vs" in line and not line.lower().startswith("date"):
            measurement_type = line
            break

    # Detect actual column header (like "V,µA" or "s,V")
    header_index = None
    for i, line in enumerate(lines):
        if ',' in line and all(any(c.isalpha() for c in token) for token in line.split(',')):
            header_index = i
            break

    if header_index is None:
        raise ValueError("Could not find a valid data header (e.g., 'V,µA' or 's,V')")

    headers = [h.strip() for h in lines[header_index].split(',')]
    data_lines = lines[header_index + 1:]

    # Read and convert numeric data only
    data = []
    for line in data_lines:
        try:
            row = [float(val) for val in line.split(',')]
            data.append(row)
        except ValueError:
            continue  # skip non-numeric lines

    df = pd.DataFrame(data, columns=headers)
    return measurement_type, df


import matplotlib.pyplot as plt

def plot_multiple_xy(datasets, labels=None, xlabel=None, ylabel=None, title=None):
    """
    Plot multiple datasets (each as a 2-column DataFrame) on one plot.

    Args:
        datasets: list of DataFrames with 2 columns
        labels: list of labels for each line
        xlabel, ylabel: optional axis labels
        title: optional plot title
    """
    fig, ax = plt.subplots(figsize=(6, 4), dpi=300)

    colors = ['blue', 'green', 'red', 'orange', 'purple']

    for i, df in enumerate(datasets):
        x = df.iloc[:, 0]
        y = df.iloc[:, 1]
        label = labels[i] if labels else None
        color = colors[i % len(colors)]
        ax.plot(x, y, label=label, color=color, linewidth=1)

    # X and Y = 0 guide lines
    ax.axhline(0, color='gray', linestyle='-', linewidth=0.5)
    ax.axvline(0, color='gray', linestyle='-', linewidth=0.5)

    # Axis style: boxed
    for side in ['top', 'bottom', 'left', 'right']:
        ax.spines[side].set_linewidth(1)
        ax.spines[side].set_color('black')

    # Axis labels
    ax.set_xlabel(xlabel if xlabel else datasets[0].columns[0], fontsize=12, fontweight='bold', labelpad=10)
    ax.set_ylabel(ylabel if ylabel else datasets[0].columns[1], fontsize=12, fontweight='bold', labelpad=10)

    # Axis limits
        # Extract x/y and lengths
    x_list = [df.iloc[:, 0].values for df in datasets]
    y_list = [df.iloc[:, 1].values for df in datasets]
    all_x = np.concatenate(x_list)
    all_y = np.concatenate(y_list)
    # ax.set_xlim(all_x.min()*1.1, all_x.max()*1.1)
    # ax.set_ylim(all_y.min()*1.2, all_y.max()*1.1)
        # Add 10% margin to x and y limits
    x_margin = (all_x.max() - all_x.min()) * 0.05
    y_margin = (all_y.max() - all_y.min()) * 0.05
    
    ax.set_xlim(all_x.min() - x_margin, all_x.max() + x_margin)
    ax.set_ylim(all_y.min() - y_margin, all_y.max() + y_margin)


    # Ticks and layout
    ax.tick_params(direction='in', top=True, right=True)
    ax.grid(False)

    # Title and legend
    if title:
        ax.set_title(title, fontsize=12,fontweight='bold', pad=10)
    if labels:
        ax.legend()

    plt.tight_layout()
    plt.show()


import matplotlib.pyplot as plt

def plot_xy(df, measurement_type=None):
    """
    Plot single xy curve
    """
    x_col = df.columns[0]
    y_col = df.columns[1]
    x = df[x_col]
    y = df[y_col]

    fig, ax = plt.subplots(figsize=(6,4),dpi=300)  # square aspect ratio

    # Plot data
    ax.plot(x, y, linewidth=1, color='blue')

    # Keep left and bottom axis
    ax.spines['left'].set_visible(True)
    ax.spines['bottom'].set_visible(True)

    # Show top and right spine to make a box
    ax.spines['top'].set_visible(True)
    ax.spines['right'].set_visible(True)

    # Draw all 4 spines in black (like a square)
    for side in ['top', 'bottom', 'left', 'right']:
        ax.spines[side].set_linewidth(1)
        ax.spines[side].set_color('black')

    # Add x=0 and y=0 guide lines
    ax.axhline(0, color='orange', linestyle='-', linewidth=0.5)
    ax.axvline(0, color='orange', linestyle='-', linewidth=0.5)

    # Axis label styling
    ax.set_xlabel('Voltage (V)', fontsize=12, fontweight='bold', labelpad=10)
    ax.set_ylabel('Current (uA)' , fontsize=12, fontweight='bold', labelpad=10)

    # Title
    if measurement_type:
        ax.set_title(measurement_type, fontsize=12, fontweight='bold', pad=10)

    # Tick styling
    ax.tick_params(direction='in', top=True, right=True)

    # Turn off grid
    ax.grid(False)

    plt.tight_layout()
    plt.show()


# Convert the plot to movie

import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import cv2
import tifffile
from matplotlib.backends.backend_agg import FigureCanvasAgg
import tqdm

def animate_CVplots(datasets, labels=["Scan 1", "Scan 2", "Scan 3"], xlabel="Voltage (V)", ylabel="Current (µA)", title="CV Scan",
                       fps=60, save_basename="cv_animation"):
    """
    Function to animate the plot
      datasets,             # List of DataFrames, each with 2 columns (x, y) for a scan
    labels=None,          # Optional list of labels for each scan (used in legend)
    xlabel=None,          # Optional label for x-axis (default: first column name)
    ylabel=None,          # Optional label for y-axis (default: second column name)
    title=None,           # Optional plot title
    fps=50,          # fps of the input data. Original speed: 4 fps
    save_basename="cv_animation"  # Prefix for saved output files (AVI + TIFF stack)

    """
     # Setup plot and canvas
    fig, ax = plt.subplots(figsize=(6, 4), dpi=150)
    fig.subplots_adjust(left=0.15, right=0.95, top=0.9, bottom=0.15)
    canvas = FigureCanvasAgg(fig)

    # Extract x/y and lengths
    x_list = [df.iloc[:, 0].values for df in datasets]
    y_list = [df.iloc[:, 1].values for df in datasets]
    lengths = [len(x) for x in x_list]
    cum_lengths = np.cumsum(lengths)
    total_frames = cum_lengths[-1]

    # Axis formatting
    for side in ['top', 'bottom', 'left', 'right']:
        ax.spines[side].set_linewidth(1)
        ax.spines[side].set_color('black')
        
    ax.axhline(0, color='gray', linewidth=0.5)
    ax.axvline(0, color='gray', linewidth=0.5)
    ax.tick_params(direction='in', top=True, right=True)
    ax.set_xlabel(xlabel or datasets[0].columns[0], fontsize=12, fontweight='bold', labelpad=10)
    ax.set_ylabel(ylabel or datasets[0].columns[1], fontsize=12, fontweight='bold', labelpad=10)
    if title:
        ax.set_title(title, fontsize=10)

    # Prepare plot elements
    colors = ['blue', 'green', 'red', 'purple']
    lines, tips = [], []
    for i in range(len(datasets)):
        line, = ax.plot([], [], color=colors[i % len(colors)], label=(labels[i] if labels else None), linewidth=1)
        tip, = ax.plot([], [], color=colors[i % len(colors)], marker='o', markersize=4)
        lines.append(line)
        tips.append(tip)

    # Add live current value in top-left
    value_text = ax.text(0.02, 0.95, '', transform=ax.transAxes,
                         fontsize=10, color='black', ha='left', va='top')


    if labels:
        ax.legend()

    # Axis limits
    all_x = np.concatenate(x_list)
    all_y = np.concatenate(y_list)

    # Add 10% margin to x and y limits
    x_margin = (all_x.max() - all_x.min()) * 0.05
    y_margin = (all_y.max() - all_y.min()) * 0.05
    
    ax.set_xlim(all_x.min() - x_margin, all_x.max() + x_margin)
    ax.set_ylim(all_y.min() - y_margin, all_y.max() + y_margin)
    # ax.set_xlim(all_x.min()*1.1, all_x.max()*1.1)
    # ax.set_ylim(all_y.min()*1.2, all_y.max()*1.1)

    # Video writer setup
    canvas.draw()
    width, height = canvas.get_width_height()
    # fps=int(1000/interval)
    aviout=save_basename+"_"+str(fps)+".avi"
    out = cv2.VideoWriter(aviout, cv2.VideoWriter_fourcc(*'MJPG'), fps, (width, height))
    # tiff_frames = []

    # Setup TIFF writer with BigTIFF support
    tif_writer = tifffile.TiffWriter(f"{save_basename}.tif", bigtiff=True)
    
    # Render frames
    print("Plotting the curves into movies~")
    for frame in tqdm.tqdm(range(total_frames)):
        ax.set_title(title, fontsize=12,fontweight='bold')  # reset title in case of matplotlib bugs

        for i, (x, y, line, tip) in enumerate(zip(x_list, y_list, lines, tips)):
            start = cum_lengths[i - 1] if i > 0 else 0
            end = cum_lengths[i]
            if frame < start:
                # Not started yet
                line.set_data([], [])
                tip.set_data([], [])
            elif frame >= end:
                # Completed
                line.set_data(x, y)
                tip.set_data([], [])
            else:
                # Active scan
                idx = frame - start + 1
                line.set_data(x[:idx], y[:idx])
                tip.set_data([x[idx - 1]], [y[idx - 1]])
                value_text.set_text(
                    f"{labels[i] if labels else f'Scan {i+1}'}\nV = {x[idx - 1]:.3f} V\nI = {y[idx - 1]:.3f} µA"
                )

        # Render and convert to image
        canvas.draw()
        buf = np.frombuffer(canvas.tostring_rgb(), dtype=np.uint8)
        image = buf.reshape((height, width, 3))
        image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        out.write(image_bgr)
        # Save TIFF frame (with compression)
        tif_writer.write(image, photometric='rgb', compression='deflate')

    out.release()
    tif_writer.close()

    # tifffile.imwrite(f"{save_basename}.tif", tiff_frames, photometric='rgb',compression='deflate')
    print("Conversion done!")
    print(f"Saved AVI: {save_basename}.avi")
    print(f"Saved TIFF stack: {save_basename}.tif")

from tkinter import Tk, filedialog
def select_multiple_files(filetypes=(("csv files", "*.csv"), ("All files", "*.*"))):
    Tk().withdraw()  # Hide the root window
    filepaths = filedialog.askopenfilenames(title="Select PSTrace .csv files", filetypes=filetypes)
    return list(filepaths)  # Convert from tuple to list
