import os 
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

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_simplecsv(filename, plot_config):
    print(f"Reading CSV file {filename}...");

    series_values = plot_config.get('load_values', ['x', 'y']);
    series_tags = plot_config.get('load_tags', None);

    plot_data = {
        'steps': []
    };

    with io.open(filename, mode='r') as file:
        column_names = file.readline().rstrip("\n").split(',');

        for line in file:
            line = line.rstrip("\n").split(',');

            line_data = {};
            for entry in enumerate(line):
                column = column_names[entry[0]];

                if(column in series_values):
                    line_data[column] = float(entry[1]);
                elif(column in (series_tags or [])):
                    line_data[column] = entry[1];

            series_name = "default"
            if(not (series_tags is None)):
                series_name = ""
                for tag in series_tags:
                    series_name = series_name + ' ' + line_data[tag];
            
            if(not (series_name in plot_data)):
                new_series = { "step": series_name };
                for value_name in series_values:
                    new_series[value_name] = [];

                plot_data[series_name] = new_series;
                plot_data['steps'].append(new_series);

            series = plot_data[series_name];

            for value_name in series_values:
                series[value_name].append(line_data[value_name]);

    return plot_data

def read_multicsv(plot_config):
    series_collection = {
        'steps': []
    }

    print("Reading multiple CSV files...")
    print(plot_config['load'])

    plot_list = plot_config['load']

    for plot_key in plot_list:
        single_plot_data = read_simplecsv(os.path.join(YAML_DIR, plot_list[plot_key]), plot_config)
        single_plot_data = single_plot_data['steps'][0]

        single_plot_data['step'] = plot_key

        series_collection[plot_key] = single_plot_data
        series_collection['steps'].append(single_plot_data)

    return series_collection

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('ymin' in plot_config):
        ax.set_ylim(bottom=plot_config['ymin']);
    if('ymax' in plot_config):
        ax.set_ylim(top=plot_config['ymax']);

    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"

    if('legend_title' in plot_config):
        legend.set_title(plot_config['legend_title']);

    ax.grid(True);

def plot_single_graph(fig, plot_config, plot_data):
    ax = fig.add_subplot();
    
    x_key = plot_config['x_key'] or 'x'
    y_key = plot_config['y_key'] or 'y'

    plot_data = plot_data['steps'][0]

    x_data = plot_data[x_key];
    y_data = plot_data[y_key];

    ax.plot(x_data, y_data, linewidth=plot_config.get('linewidth'));

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

    decorate_ax(ax, plot_config)

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'])
    
    for idx, step in enumerate(plot_data['steps']):
        plot_params = dict()
        plot_params['label'] = step['step']

        if(plot_config.get('colourmap', 'coolwarm') == 'coolwarm'):
            cmap = plt.cm.coolwarm
            plot_params['color'] = cmap(idx/(num_steps-1))
        
        plot_params['linewidth'] = plot_config.get('linewidth');

        ax.plot(step[x_key], step[y_key], **plot_params)

    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 perform_bandwidth_normalization(plot_data, plot_config):
    print("Normalizing bandwidth for all steps...")

    for step in plot_data['steps']:
        y_data = step[plot_config['y_key']]

        y_adjust = y_data[0] - plot_config.get('bandwidth_zero', 0)

        new_y_data = []

        for datapoint in y_data:
            new_y_data.append(datapoint - y_adjust)

        step[plot_config['y_key']] = new_y_data

def perform_peak_normalization(plot_data, plot_config):
    print("Normalizing peak height to 1")

    for step in plot_data['steps']:
        y_data = step[plot_config['y_key']]
        new_y_data = []

        y_max = max(y_data)
        y_min = min(y_data)

        scaling_factor = 0.1 + 0.9*(y_max if (y_max > (-y_min)) else y_min)

        for datapoint in y_data:
            new_y_data.append(datapoint / scaling_factor)

        step[plot_config['y_key']] = new_y_data

def perform_offset_removal(plot_data, plot_config):
    print("Removing offset")

    for step in plot_data['steps']:
        y_data = step[plot_config['y_key']]
        new_y_data = []

        offset_value = np.percentile(y_data, 30) * 0.8

        for datapoint in y_data:
            new_y_data.append(datapoint - offset_value)

        step[plot_config['y_key']] = new_y_data


def perform_processing_step(data_process_step, plot_data, plot_config):
    if(data_process_step == 'normalize_bandwidth'):
        perform_bandwidth_normalization(plot_data, plot_config)
    if(data_process_step == 'remove_offset'):
        perform_offset_removal(plot_data, plot_config)
    if(data_process_step == 'normalize_peak'):
        perform_peak_normalization(plot_data, 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);

        if(plot_config['loadtype'] == 'simplecsv'):
            plot_data = read_simplecsv(os.path.join(YAML_DIR, plot_config['load']), plot_config);

        if(plot_config['loadtype'] == 'multicsv'):
            plot_data = read_multicsv(plot_config)

    for data_process_step in plot_config.get('data_processing_steps', []):
        perform_processing_step(data_process_step, plot_data, plot_config)

    fig = plt.figure(figsize=(6.5, 4));

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

    if(plot_config['type'] == 'single'):
        plot_single_graph(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));

    matplotlib.pyplot.close()
    


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