#!python
# Copyright 2020,2021 Sony Corporation.
# Copyright 2021,2022,2023 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import pathlib

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


def smooth_with_moving_average(data, k=1):
    if k == 1:
        # no smoothing
        return data
    # assuming that the data is flat
    data = np.insert(data, obj=0, values=[data[0]] * (k - 1))
    cumsum = np.cumsum(data, dtype=float)
    cumsum[k:] = cumsum[k:] - cumsum[:-k]
    moving_average = cumsum[k - 1:] / k
    return moving_average


def load_data(path):
    path = pathlib.Path(path)
    return np.loadtxt(str(path), delimiter='\t', skiprows=1)


def plot_results(args):
    plt.style.use('tableau-colorblind10')
    fig = plt.figure(figsize=(5, 4), dpi=80)
    ax = fig.add_subplot(111)

    xnlim = args.xnlim
    xplim = args.xplim
    ynlim = args.ynlim
    yplim = args.yplim

    for i, tsvpath in enumerate(args.tsvpaths):
        results = load_data(tsvpath)

        itr_x = results[:, 0]
        avg_y = results[:, 1]
        std_y = results[:, 2]
        min_y = results[:, 3]
        max_y = results[:, 4]
        med_y = results[:, 5]

        x = itr_x
        if args.xnlim is None:
            xnlim = np.amin(x)
        else:
            xnlim = min(np.amin(x), xnlim)
        if args.xplim is None:
            xplim = np.amax(x)
        else:
            xplim = max(np.amax(x), xplim)
        if args.ynlim is None:
            ynlim = np.amin(min_y)
        else:
            ynlim = max(np.amin(min_y), ynlim)
        if args.yplim is None:
            yplim = np.amax(max_y)
        else:
            yplim = max(np.amax(max_y), yplim)

        mean_label = 'mean'
        if i < len(args.tsvlabels):
            mean_label = f'{mean_label}({args.tsvlabels[i]})'
        avg_y = smooth_with_moving_average(avg_y, k=args.smooth_k)
        ax.plot(x, avg_y, label=mean_label, linewidth=1)
        if not args.no_stddev:
            ax.fill_between(x, avg_y + std_y, avg_y - std_y, alpha=0.3)

        if args.plot_median:
            median_label = 'median'
            if i < len(args.tsvlabels):
                median_label = f'{median_label}({args.tsvlabels[i]})'
                med_y = smooth_with_moving_average(med_y, k=args.smooth_k)
                ax.plot(x, med_y, label=median_label, linewidth=1)
    ax.set_xlim(args.xnlim, args.xplim)
    ax.set_ylim(args.ynlim, args.yplim)
    ax.set_xlabel(args.xlabel)
    ax.set_ylabel(args.ylabel)
    ax.xaxis.set_major_formatter(ScalarFormatter(useMathText=True))
    ax.ticklabel_format(style="sci",  axis="x", scilimits=(0, 0))
    for height in args.hlines:
        plt.axhline(y=height, color='red')

    ax.legend(loc=args.legend_pos, fontsize=8)

    plt.grid(True)
    plt.tight_layout()
    plt.savefig(args.outdir + '/' + args.savename)
    if args.show_fig:
        plt.show()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--tsvpaths',
                        type=str,
                        nargs='+',
                        required=True,
                        help='evaluation tsv filepath')
    parser.add_argument('--tsvlabels',
                        type=str,
                        nargs='+',
                        default=[],
                        help='evaluation tsv labels')

    # Graph options
    parser.add_argument('--xlabel', type=str, default='steps')
    parser.add_argument('--ylabel', type=str, default='score')

    parser.add_argument('--outdir', type=str, default='./')
    parser.add_argument('--savename', type=str, default='result.png')

    parser.add_argument('--xnlim', type=int, default=None)
    parser.add_argument('--xplim', type=int, default=None)
    parser.add_argument('--ynlim', type=int, default=None)
    parser.add_argument('--yplim', type=int, default=None)

    parser.add_argument('--legend-pos', type=str, default='upper left')

    parser.add_argument('--hlines', type=float, nargs='*', default=[])

    parser.add_argument('--plot-median', action='store_true')
    parser.add_argument('--no-stddev', action='store_true')
    parser.add_argument('--smooth-k', type=int, default=1)

    parser.add_argument('--show-fig', action='store_true')

    args = parser.parse_args()

    plot_results(args)
