import sys
import io
import re

import yaml

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.ticker import EngFormatter

import colorsys

import os 
SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))

class Re(object):
  def __init__(self):
    self.last_match = None
  def match(self,pattern,text):
    self.last_match = re.match(pattern,text)
    return self.last_match
  def search(self,pattern,text):
    self.last_match = re.search(pattern,text)
    return self.last_match


def setup_step(plot_data, key):
    if(key == None):
        key = "default";
    
    if(key in plot_data):
        return;

    next_step = {
        "step" : key
    };

    plot_data[key] = next_step;
    plot_data['steps'].append(next_step);

def read_ltspice_file(filename, plot_config):
    print(f"Reading LTSpice .txt file {filename}...");

    series = [];
    plot_data  = {
        'steps': []
    };
    current_step = None;

    with io.open(filename, mode="r", encoding="ISO8859") as file:
        header = file.readline();
        lines = header.rstrip("\n").split("\t");
        
        for line in file:
            line = line.rstrip("\n");
            lre = Re();
            
            if lre.search("^Step Information: ([^=]*)=([\d\.\+-]+)(\w?)  \(Step: .*\)$", line):
                step_param = plot_config.get('step_parameter', lre.last_match[1]);

                value = f"{round(float(lre.last_match[2]), 2):.2f}".replace('.', ',');
                unit = f"{lre.last_match[3]}{plot_config.get('step_unit', '')}";

                current_step = "{value:>6s} {unit:s}".format(value=value, unit=unit);
            else:
                setup_step(plot_data, current_step);
                step = plot_data[current_step];
                
                elements = line.split("\t");
                for idx, element in enumerate(elements):
                    lname = lines[idx];
                    ere = Re();
                    if ere.search("^([\d\.e\+-]+)$", element):
                        if(not lname in step):
                            step[lname] = []
                        step[lname].append(float(ere.last_match[1]));
                    elif ere.search("^\(([\d\.e\+-]+)dB,([\d\.e\+-]+)", element):
                        if(not (lname+" dB") in step):
                            step[lname+" dB"] = []
                            step[lname+" deg"] = []
                        step[lname+" dB"].append(float(ere.last_match[1]));
                        step[lname+" deg"].append(float(ere.last_match[2]));
                    else:
                        raise RuntimeError("Unknown/Not configured parsing element!");
    
    return plot_data;

def decorate_ax(ax, plot_config):
    if(plot_config.get('show_title', False)):
        ax.set_title(plot_config['title']);

    ax.set_xlabel(plot_config['xlabel']);
    ax.set_ylabel(plot_config['ylabel']);
    
    if('yscale' in plot_config):
        ax.set_yscale(plot_config['yscale']);
    if('xscale' in plot_config):
        ax.set_xscale(plot_config['xscale']);

    if('xmin' in plot_config):
        ax.set_xlim(left=plot_config['xmin']);
    if('xmax' in plot_config):
        ax.set_xlim(right=plot_config['xmax']);

    if('xformatter' in plot_config):
        if('engineering' == plot_config['xformatter']):
            formatter = EngFormatter(places=plot_config.get('xplaces', 0), sep="\N{THIN SPACE}")
            ax.xaxis.set_major_formatter(formatter)

    if('yformatter' in plot_config):
        if('engineering' == plot_config['yformatter']):
            formatter = EngFormatter(places=plot_config.get('yplaces', 0))
            ax.yaxis.set_major_formatter(formatter)

    legend = ax.legend();
    hp = legend._legend_box.get_children()[1]
    for vp in hp.get_children():
        for row in vp.get_children():
            row.set_width(350)  # need to adapt this manually
            row.mode= "expand"
            row.align="right"

    ax.grid(True);

def plot_lt_sweep(fig, plot_config, plot_data):
    step_keys = plot_data.keys();

    ax = fig.add_subplot();
    ax.set_xscale('log');

    x_key = "Freq.";
    if("x_key" in plot_config):
        x_key = plot_config['x_key'];
    
    y_key = None;
    if('y_key' in plot_config):
        y_key = plot_config['y_key'];
    
    if(y_key == None):
        raise RuntimeError("No Y-Data Key (`y_key`) specified for plot!");

    num_steps = len(plot_data['steps']);
    cmap = plt.cm.coolwarm;

    custom_lines = [matplotlib.lines.Line2D([0], [0], color=cmap(0.), lw=4),
                matplotlib.lines.Line2D([0], [0], color=cmap(.5), lw=4),
                matplotlib.lines.Line2D([0], [0], color=cmap(1.), lw=4)];

    for idx, step in enumerate(plot_data['steps']):
        ax.plot(step[x_key], step[y_key], color=cmap(idx/(num_steps-1)), label=step['step']);

    if(not 'xformatter' in plot_config):
        plot_config['xformatter'] = 'engineering';

    if(not 'xmin' in plot_config):
        plot_config['xmin'] = np.min(plot_data['steps'][0][x_key]);
    if(not 'xmax' in plot_config):
        plot_config['xmax'] = np.max(plot_data['steps'][0][x_key]);

    ax.legend();

    decorate_ax(ax, plot_config);

def generate_plot(plot_config):
    global YAML_DIR;

    plot_data = None;
    if("load" in plot_config):
        if(not "loadtype" in plot_config):
            raise RuntimeError("Missing load type (`loadtype`) for plot config");
        
        if(plot_config['loadtype'] == 'ltspice'):
            plot_data = read_ltspice_file(os.path.join(YAML_DIR, plot_config['load']), plot_config);

    fig = plt.figure();

    if(plot_config['type'] == 'lt_sweep'):
        plot_lt_sweep(fig, plot_config, plot_data);

    fig.subplots_adjust(0.15, 0.12, 0.96, 0.9)

    fig.savefig(os.path.join(YAML_DIR, plot_config['ofile']), dpi=plot_config.get('dpi', 300));
    


INPUT_YAML_FILE = SCRIPT_DIR + "/plots.yml" if (len(sys.argv) <= 1) else sys.argv[1];
YAML_DIR = os.path.dirname(INPUT_YAML_FILE);
PLOT_CONFIG = None;

print(f"Reading YAML config {INPUT_YAML_FILE}");

with open(INPUT_YAML_FILE, "r") as file:
    PLOT_CONFIG = yaml.load(file, yaml.Loader);


for plot in PLOT_CONFIG['plots']:
    plot = {**PLOT_CONFIG['defaults'], **plot};

    generate_plot(plot);