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):
    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: (.*)  \(Step: .*\)$", line):
                current_step = lre.last_match[1];
            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):
    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('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)

    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)));

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

    legend_data = [];
    for x in [0, int(num_steps/2), num_steps-1]:
        step_name = plot_data['steps'][x]['step'];
        for orig, replacement in plot_config.get('legend_replace', dict()).items():
            step_name = step_name.replace(orig, replacement);
        
        legend_data.append(step_name);

    ax.legend(custom_lines, legend_data);

    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']));

    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);