################################################################################
#
# Copyright (C) 2020 Marvell Technology Group Ltd.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
# DAMAGE.
#
################################################################################

################################################################################
# This script requires next python modules:
#
#     sudo python -m pip install matplotlib
#
################################################################################

import argparse
import json
import matplotlib.pyplot as plt
import matplotlib.colors as clr
import matplotlib.ticker as tck
import matplotlib.gridspec as gsp
import numpy


def get_color_map():
    # PLOT COLOR MAP =====================================================
    r = 0; g = 1; b = 2
    # --------------------------------------------------------------------
    map_color = numpy.zeros((512, 3))
    # --------------------------------------------------------------------
    map_color[0, r] = 1
    map_color[0, g] = 1
    map_color[0, b] = 1
    # --------------------------------------------------------------------
    map_color[1, r] = 0.6
    map_color[1, g] = 0.6
    map_color[1, b] = 0.6
    # --------------------------------------------------------------------
    map_color[2, r] = 0.5
    map_color[2, g] = 0.5
    map_color[2, b] = 0.5
    # --------------------------------------------------------------------
    map_color[3:8, r] = 0.4
    map_color[3:8, g] = 0.4
    map_color[3:8, b] = 0.4
    # --------------------------------------------------------------------
    map_color[8:16, r] = 0.3
    map_color[8:16, g] = 0.3
    map_color[8:16, b] = 0.3
    # --------------------------------------------------------------------
    map_color[16:32, r] = 0.2
    map_color[16:32, g] = 0.2
    map_color[16:32, b] = 0.2
    # --------------------------------------------------------------------
    map_color[32:65, r] = 1
    map_color[32:65, g] = 0
    map_color[32:65, b] = 0
    # --------------------------------------------------------------------
    map_color[64:128, r] = 1
    map_color[64:128, g] = 0.5
    map_color[64:128, b] = 0
    # --------------------------------------------------------------------
    map_color[128:256, r] = 1
    map_color[128:256, g] = 1
    map_color[128:256, b] = 0
    # --------------------------------------------------------------------
    map_color[255:512, :] = 1
    map_color[255:512, g] = 1
    map_color[255:512, b] = 1
    # ====================================================================

    return clr.ListedColormap(map_color)


def prepare_eye_data(m_eye_vals, pttrn_list, byp_rx):
    num_phs_steps = numpy.size(m_eye_vals, 2)
    num_dac_steps = numpy.size(m_eye_vals, 1)

    d_eye_vals = numpy.zeros((numpy.size(pttrn_list, 0), num_dac_steps - 1, num_phs_steps))

    for pttrn_iter in range(len(pttrn_list)):
        for phs_iter in range(num_phs_steps):
            # CALCULATE DIFFERENCE ---------------------------------------
            d_eye_vals[pttrn_iter, :, phs_iter] = numpy.abs(numpy.diff(m_eye_vals[pttrn_iter, :, phs_iter]))
            # ------------------------------------------------------------

    # COMBINE RESULTS ----------------------------------------------------
    xd_eye_vals = numpy.squeeze(numpy.sum(d_eye_vals, 0))
    # --------------------------------------------------------------------
    xd_eye_vals[1:-1, :] = xd_eye_vals[1:-1, :] + xd_eye_vals[2:, :] + xd_eye_vals[:-2, :]
    # --------------------------------------------------------------------
    xd_eye_vals = numpy.abs(xd_eye_vals / numpy.max(numpy.max(xd_eye_vals))) * 512
    # --------------------------------------------------------------------

    if not byp_rx and len(pttrn_list) > 2:
        # COMBINE RESULTS ------------------------------------------------
        xd_eye_vals_1x = numpy.squeeze(d_eye_vals[numpy.where(pttrn_list == 6), :, :]) + numpy.squeeze(d_eye_vals[numpy.where(pttrn_list == 4), :, :])
        # ----------------------------------------------------------------
        xd_eye_vals_1x[1:-1, :] = xd_eye_vals_1x[1:-1, :] + xd_eye_vals_1x[2:, :] + xd_eye_vals_1x[:-2, :]
        # ----------------------------------------------------------------
        xd_eye_vals_1x = numpy.abs(xd_eye_vals_1x / numpy.max(numpy.max(xd_eye_vals_1x))) * 512
        # ----------------------------------------------------------------

        # COMBINE RESULTS ------------------------------------------------
        xd_eye_vals_0x = numpy.squeeze(d_eye_vals[numpy.where(pttrn_list == 2), :, :]) + numpy.squeeze(d_eye_vals[numpy.where(pttrn_list == 0), :, :])
        # ----------------------------------------------------------------
        xd_eye_vals_0x[1:-1, :] = xd_eye_vals_0x[1:-1, :] + xd_eye_vals_0x[2:, :] + xd_eye_vals_0x[:-2, :]
        # ----------------------------------------------------------------
        xd_eye_vals_0x = numpy.abs(xd_eye_vals_0x / numpy.max(numpy.max(xd_eye_vals_0x))) * 512
        # ----------------------------------------------------------------
    else:
        xd_eye_vals_1x = numpy.zeros(numpy.shape(xd_eye_vals))
        xd_eye_vals_0x = numpy.zeros(numpy.shape(xd_eye_vals))

    return xd_eye_vals_1x, xd_eye_vals_0x, xd_eye_vals


def calc_eye_dimensions(xd_eye_vals):
    num_phs_steps = numpy.size(xd_eye_vals, 1)
    num_dac_steps = numpy.size(xd_eye_vals, 0)

    phs_center = num_phs_steps * 2 // 3
    dac_center = num_dac_steps // 2

    eye_width = 0

    for phs_iter in range(phs_center + 1, num_phs_steps):
        if int(xd_eye_vals[dac_center][phs_iter]) != 0:
            break
        eye_width += 1

    for phs_iter in range(phs_center, -1, -1):
        if int(xd_eye_vals[dac_center][phs_iter]) != 0:
            break
        eye_width += 1

    eye_height = 0

    for dac_iter in range(dac_center + 1, num_dac_steps):
        if int(xd_eye_vals[dac_iter][phs_center]) != 0:
            break
        eye_height += 1

    for dac_iter in range(dac_center, -1, -1):
        if int(xd_eye_vals[dac_iter][phs_center]) != 0:
            break
        eye_height += 1

    eye_width_ui = (1.5 / num_phs_steps) * eye_width
    eye_height_mv = (650.0 / num_dac_steps) * eye_height

    return eye_width_ui, eye_height_mv


def get_rate_rx_str(rate_rx):
    if rate_rx == 0:
        return "Gen4 (16 GT/s)"
    if rate_rx == 1:
        return "Gen3 (8 GT/s)"
    if rate_rx == 2:
        return "Gen2 (5 GT/s)"
    if rate_rx == 3:
        return "Gen1 (2.5 GT/s)"
    return "Unknown ({})".format(rate_rx)


if __name__ == "__main__":
    # PARSE COMMAND LINE ARGUMENTS =======================================
    parser = argparse.ArgumentParser()
    parser.add_argument("json")
    args = parser.parse_args()

    file_name = args.json
    # ====================================================================

    # LOAD RAW DATA ======================================================
    with open(file_name, "r") as fp:
        scope_results = json.load(fp)

    phs_step = scope_results["phs_step"]
    dac_step = scope_results["dac_step"]
    data_msk = scope_results["data_msk"]
    stat_len = scope_results["stat_len"]

    rate_rx = scope_results["rate_rx"]
    byp_rx = scope_results["byp_rx"]
    pttrn_list = numpy.array(scope_results["pttrn_list"])
    m_eye_vals = numpy.array(scope_results["m_eye_vals"])
    # ====================================================================

    # PLOT COLOR MAP =====================================================
    map_color = get_color_map()
    # ====================================================================

    # SET RANGES ---------------------------------------------------------
    x_axis_min = -1.0
    x_axis_max = 0.5
    y_axis_min = -325.0
    y_axis_max = 325.0
    # --------------------------------------------------------------------
    xd_eye_vals_1x, xd_eye_vals_0x, xd_eye_vals = prepare_eye_data(m_eye_vals, pttrn_list, byp_rx)
    # --------------------------------------------------------------------
    eye_width, eye_height = calc_eye_dimensions(xd_eye_vals)
    # --------------------------------------------------------------------

    imshow_kwargs = {
        "aspect": (x_axis_max - x_axis_min) / (y_axis_max - y_axis_min),
        "extent": [x_axis_min, x_axis_max, y_axis_min, y_axis_max],
        "origin": "lower",
        "interpolation": "none",
        "cmap": map_color
    }

    fig = plt.figure(figsize=(8, 7), constrained_layout=True)
    gs = gsp.GridSpec(2, 2, figure=fig)
    axes = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]), fig.add_subplot(gs[1, 0])]

    for i, (data, title) in enumerate(zip((xd_eye_vals_1x, xd_eye_vals_0x, xd_eye_vals),
                                          ("1X eye", "0X eye", "effective eye"))):
        # PLOT EYE -------------------------------------------------------
        im = axes[i].imshow(data, **imshow_kwargs)
        axes[i].set_title(title)
        axes[i].set_xlabel("Time (UI)")
        axes[i].set_ylabel("Amplitude (mV)")
        axes[i].xaxis.set_major_locator(tck.MultipleLocator(0.5))
        axes[i].yaxis.set_major_locator(tck.MultipleLocator(50))
        axes[i].grid(True)
        # ----------------------------------------------------------------

    # SET PLOT ATTRIBUTES ------------------------------------------------
    fig.colorbar(im, ax=axes, shrink=0.7, aspect=30)
    # ----------------------------------------------------------------
    window_title = "SerDes Eye Plot - {}".format(file_name)
    fig.canvas.set_window_title(window_title)
    fig.suptitle(window_title)
    # ----------------------------------------------------------------
    fig.text(0.55, 0.4,
             "Scope Parameters:\n"
             "    phs_step: {phs_step}\n"
             "    dac_step: {dac_step}\n"
             "    data_msk: {data_msk}\n"
             "    stat_len: {stat_len}\n"
             "\n"
             "Device Parameters:\n"
             "    Rate RX: {rate_rx}\n"
             "    Eye Width: {eye_width:.4f} UI\n"
             "    Eye Height: {eye_height:.4f} mV\n".format(
                 phs_step=phs_step, dac_step=dac_step, data_msk=data_msk, stat_len=stat_len,
                 rate_rx=get_rate_rx_str(rate_rx), eye_width=eye_width, eye_height=eye_height),
             verticalalignment="top")
    # ----------------------------------------------------------------

    plt.show()
