import argparse

import numpy as np
import os
import sys

import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt

sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from saliency.visualizer.smiles_visualizer import SmilesVisualizer


def visualize(dir_path):
    parent_dir = os.path.dirname(dir_path)
    saliency_vanilla = np.load(os.path.join(dir_path, "saliency_vanilla.npy"))
    saliency_smooth = np.load(os.path.join(dir_path, "saliency_smooth.npy"))
    saliency_bayes = np.load(os.path.join(dir_path, "saliency_bayes.npy"))

    visualizer = SmilesVisualizer()
    os.makedirs(os.path.join(parent_dir, "result_vanilla"), exist_ok=True)
    os.makedirs(os.path.join(parent_dir, "result_smooth"), exist_ok=True)
    os.makedirs(os.path.join(parent_dir, "result_bayes"), exist_ok=True)

    test_idx = np.load(os.path.join(dir_path, "test_idx.npy"))
    answer = np.load(os.path.join(dir_path, "answer.npy"))
    output = np.load(os.path.join(dir_path, "output.npy"))

    smiles_all = np.load(os.path.join(parent_dir, "smiles.npy"))

    def calc_range(saliency):
        vmax = float('-inf')
        vmin = float('inf')
        for v in saliency:
            vmax = max(vmax, np.max(v))
            vmin = min(vmin, np.min(v))
        return vmin, vmax

    v_range_vanilla = calc_range(saliency_vanilla)
    v_range_smooth = calc_range(saliency_smooth)
    v_range_bayes = calc_range(saliency_bayes)

    def get_scaler(v_range):
        def scaler(saliency_):
            saliency = np.copy(saliency_)
            minv, maxv = v_range
            if maxv == minv:
                saliency = np.zeros_like(saliency)
            else:
                pos = saliency >= 0.0
                saliency[pos] = saliency[pos]/maxv
                nega = saliency < 0.0
                saliency[nega] = saliency[nega]/(np.abs(minv))
            return saliency
        return scaler

    scaler_vanilla = get_scaler(v_range_vanilla)
    scaler_smooth = get_scaler(v_range_smooth)
    scaler_bayes = get_scaler(v_range_bayes)

    def color(x):
        if x > 0:
            # Red for positive value
            return 1., 1. - x, 1. - x
        else:
            # Blue for negative value
            x *= -1
            return 1. - x, 1. - x, 1.

    for i, id in enumerate(test_idx):
        smiles = smiles_all[id]
        out = output[i]
        ans = answer[i]
        # legend = "t:{}, p:{}".format(ans, out)
        legend = ''
        ext = '.png'  # '.svg'
        # visualizer.visualize(
        #     saliency_vanilla[id], smiles, save_filepath=os.path.join(parent_dir, "result_vanilla", str(id) + ext),
        #     visualize_ratio=1.0, legend=legend, scaler=scaler_vanilla, color_fn=color)
        # visualizer.visualize(
        #     saliency_smooth[id], smiles, save_filepath=os.path.join(parent_dir, "result_smooth", str(id) + ext),
        #     visualize_ratio=1.0, legend=legend, scaler=scaler_smooth, color_fn=color)
        visualizer.visualize(
            saliency_bayes[id], smiles, save_filepath=os.path.join(parent_dir, "result_bayes", str(id) + ext),
            visualize_ratio=1.0, legend=legend, scaler=scaler_bayes, color_fn=color)


def plot_result(prediction, answer, save_filepath='result.png'):
    plt.scatter(prediction, answer, marker='.')
    plt.plot([-100, 100], [-100, 100], c='r')
    max_v = max(np.max(prediction), np.max(answer))
    min_v = min(np.min(prediction), np.min(answer))
    plt.xlim([min_v-0.1, max_v+0.1])
    plt.xlabel("prediction")
    plt.ylim([min_v-0.1, max_v+0.1])
    plt.ylabel("ground truth")
    plt.savefig(save_filepath)
    plt.close()


def main():
    parser = argparse.ArgumentParser(
        description='Regression with own dataset.')
    parser.add_argument('--dirpath', '-d', type=str, default='./results/M_30_3_32_32')
    args = parser.parse_args()
    path = args.dirpath
    n_split = 5
    output = []
    answer = []
    for i in range(n_split):
        suffix = str(i) + "-" + str(n_split)
        output.append(np.load(os.path.join(path, suffix, "output.npy")))
        answer.append(np.load(os.path.join(path, suffix, "answer.npy")))
    output = np.concatenate(output)
    answer = np.concatenate(answer)

    plot_result(output, answer, save_filepath=os.path.join(path, "result.png"))
    for i in range(n_split):
        suffix = str(i) + "-" + str(n_split)
        print(suffix)
        visualize(os.path.join(path, suffix))


if __name__ == '__main__':
    main()


[(8, 'matplotlib.use', 'matplotlib.use', (['"""agg"""'], {}), False, 'import matplotlib\n'), (16, 'os.path.dirname', 'os.path.dirname', (['dir_path'], {}), False, 'import os\n'), (21, 'saliency.visualizer.smiles_visualizer.SmilesVisualizer', 'SmilesVisualizer', ([], {}), False, 'from saliency.visualizer.smiles_visualizer import SmilesVisualizer\n'), (90, 'matplotlib.pyplot.scatter', 'plt.scatter', (['prediction', 'answer'], {'marker': '"""."""'}), True, 'import matplotlib.pyplot as plt\n'), (91, 'matplotlib.pyplot.plot', 'plt.plot', (['[-100, 100]', '[-100, 100]'], {'c': '"""r"""'}), True, 'import matplotlib.pyplot as plt\n'), (94, 'matplotlib.pyplot.xlim', 'plt.xlim', (['[min_v - 0.1, max_v + 0.1]'], {}), True, 'import matplotlib.pyplot as plt\n'), (95, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""prediction"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (96, 'matplotlib.pyplot.ylim', 'plt.ylim', (['[min_v - 0.1, max_v + 0.1]'], {}), True, 'import matplotlib.pyplot as plt\n'), (97, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""ground truth"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (98, 'matplotlib.pyplot.savefig', 'plt.savefig', (['save_filepath'], {}), True, 'import matplotlib.pyplot as plt\n'), (99, 'matplotlib.pyplot.close', 'plt.close', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (103, 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""Regression with own dataset."""'}), False, 'import argparse\n'), (115, 'numpy.concatenate', 'np.concatenate', (['output'], {}), True, 'import numpy as np\n'), (116, 'numpy.concatenate', 'np.concatenate', (['answer'], {}), True, 'import numpy as np\n'), (17, 'os.path.join', 'os.path.join', (['dir_path', '"""saliency_vanilla.npy"""'], {}), False, 'import os\n'), (18, 'os.path.join', 'os.path.join', (['dir_path', '"""saliency_smooth.npy"""'], {}), False, 'import os\n'), (19, 'os.path.join', 'os.path.join', (['dir_path', '"""saliency_bayes.npy"""'], {}), False, 'import os\n'), (22, 'os.path.join', 'os.path.join', (['parent_dir', '"""result_vanilla"""'], {}), False, 'import os\n'), (23, 'os.path.join', 'os.path.join', (['parent_dir', '"""result_smooth"""'], {}), False, 'import os\n'), (24, 'os.path.join', 'os.path.join', (['parent_dir', '"""result_bayes"""'], {}), False, 'import os\n'), (26, 'os.path.join', 'os.path.join', (['dir_path', '"""test_idx.npy"""'], {}), False, 'import os\n'), (27, 'os.path.join', 'os.path.join', (['dir_path', '"""answer.npy"""'], {}), False, 'import os\n'), (28, 'os.path.join', 'os.path.join', (['dir_path', '"""output.npy"""'], {}), False, 'import os\n'), (30, 'os.path.join', 'os.path.join', (['parent_dir', '"""smiles.npy"""'], {}), False, 'import os\n'), (92, 'numpy.max', 'np.max', (['prediction'], {}), True, 'import numpy as np\n'), (92, 'numpy.max', 'np.max', (['answer'], {}), True, 'import numpy as np\n'), (93, 'numpy.min', 'np.min', (['prediction'], {}), True, 'import numpy as np\n'), (93, 'numpy.min', 'np.min', (['answer'], {}), True, 'import numpy as np\n'), (46, 'numpy.copy', 'np.copy', (['saliency_'], {}), True, 'import numpy as np\n'), (118, 'os.path.join', 'os.path.join', (['path', '"""result.png"""'], {}), False, 'import os\n'), (122, 'os.path.join', 'os.path.join', (['path', 'suffix'], {}), False, 'import os\n'), (11, 'os.path.abspath', 'os.path.abspath', (['__file__'], {}), False, 'import os\n'), (36, 'numpy.max', 'np.max', (['v'], {}), True, 'import numpy as np\n'), (37, 'numpy.min', 'np.min', (['v'], {}), True, 'import numpy as np\n'), (49, 'numpy.zeros_like', 'np.zeros_like', (['saliency'], {}), True, 'import numpy as np\n'), (113, 'os.path.join', 'os.path.join', (['path', 'suffix', '"""output.npy"""'], {}), False, 'import os\n'), (114, 'os.path.join', 'os.path.join', (['path', 'suffix', '"""answer.npy"""'], {}), False, 'import os\n'), (54, 'numpy.abs', 'np.abs', (['minv'], {}), True, 'import numpy as np\n')]