from termcolor import colored
import random
from tqdm import tqdm
import aiohttp
import asyncio
from aiohttp import ClientSession
from helper.constants import TIME_FILTERS, URL_OPTIONS, RANDOM_TERMS
from helper.utils import get_url_option, get_time_filter, get_platform_option, get_video_option, send_discord_message
from fetcher.youtube_fetcher import (fetch_youtube_titles, fetch_video_content,
                                     extract_likes_from_content, extract_description_from_content,
                                     download_video, async_fetch_youtube_titles)
from model.model import YouTubeTitleGenerator

def get_multiple_search_queries():
    queries = []
    while True:
        query = input(colored('Enter a YouTube search query or hit ENTER to start scraping: ', 'cyan')).strip()
        if not query:
            break
        queries.append(query)
    return queries

def main():
    platform_option = get_platform_option()

    if platform_option == '1':
        url_option = get_url_option()
        if url_option == URL_OPTIONS['1']:
            query = input(colored('Enter your YouTube search query: ', 'cyan'))
            time_filter = get_time_filter()
            videos = fetch_youtube_titles(URL_OPTIONS['1'] + query, time_filter)
        else:
            videos = fetch_youtube_titles(url_option)

        sorted_videos = sorted(videos, key=lambda x: x[1], reverse=True)[:10]

        print(colored('Top 10 YouTube Videos by Views:', 'blue'))
        for idx, (title, views, link) in enumerate(sorted_videos, 1):
            print(colored(f'{idx}. Title: {title}', 'green'))
            print(colored(f'   Views: {views:,}', 'yellow'))
            print(colored(f'   Link: {link}', 'cyan'))

    elif platform_option == '2':
        video_url = input(colored('Enter the YouTube video link: ', 'cyan'))
        video_option = get_video_option()

        # Fetch and save content of the provided video URL
        fetch_video_content(video_url)

        if video_option == '1':
            likes = extract_likes_from_content()
            print(colored(f'Likes: {likes}', 'green'))
        elif video_option == '2':
            description = extract_description_from_content()
            print(colored(f'Description: {description}', 'green'))
        elif video_option == '3':
            download_video(video_url)

    elif platform_option == '3':
        queries = get_multiple_search_queries()
        with open('scraped_titles.txt', 'a', encoding='utf-8') as file:
            for query in queries:
                videos = fetch_youtube_titles(URL_OPTIONS['1'] + query)
                for title, _, _ in videos:
                    file.write(title + '\n')
    elif platform_option == '4':
        titles_to_write = []

        async def fetch_all_titles():
            async with ClientSession() as session:
                tasks = []
                for query in RANDOM_TERMS:
                    for time_key, time_filter in TIME_FILTERS.items():
                        tasks.append(async_fetch_youtube_titles(session, URL_OPTIONS['1'] + query + time_filter))
                return await asyncio.gather(*tasks)

        loop = asyncio.get_event_loop()
        all_videos = loop.run_until_complete(fetch_all_titles())

        for videos in all_videos:
            for title, _, _ in videos:
                titles_to_write.append(title)

        with open('scraped_titles.txt', 'a', encoding='utf-8') as file:
            for title in titles_to_write:
                file.write(title + '\n')
        send_discord_message("Finished scraping titles!")
    elif platform_option == '5':
        title_generator = YouTubeTitleGenerator()
        with open("scraped_titles.txt", 'r', encoding='utf-8') as file:
            titles = file.readlines()
        title_generator.train(titles)
        title_generator.save_model("model/youtube_title_generator_model.h5")
        send_discord_message("Finished training model!")

        print("Model trained and saved!")

    elif platform_option == '6':
        title_generator = YouTubeTitleGenerator()
        max_sequence_len, total_words = title_generator.load_model("model/youtube_title_generator_model.h5", titles)
        prompt = input("Enter a prompt: ")
        generated_title = title_generator.generate_title(prompt, max_sequence_len, total_words)
        print("Generated Title:", generated_title)

if __name__ == "__main__":
    main()


import os
import glob
from collections import defaultdict
from multiprocessing.pool import Pool

import numpy as np
import matplotlib
import re
import sys
import pickle as pkl

matplotlib.use('Agg')
from matplotlib import pyplot as plt
from datetime import datetime
from cycler import cycler
import xarray as xr
from utils.data import DataUtils

type_color_mapping = {
    'train': 'g',
    'test': 'c',
    'filterd': 'y'
}

# generate a dictionary {station_id -> formatted_staion_name}
def get_station_dict(observation_data, station_ids):
    station_name_dict = {}
    for station_id in station_ids:
        station_name_dict[station_id] = str(observation_data['name'].sel(station_id=station_id).data, encoding='utf-8')
    return station_name_dict

def plot_datetime_distribution(train_data, test_data, output_path, seed=None):
    # prepare output folder
    if not output_path.endswith('/'):
        output_path += '/'
    output_path += 'datetime_distribution_plots'

    if not os.path.exists(output_path):
        os.makedirs(output_path)

    # loop over folds
    for fold in range(len(train_data)):
        # sort train and test data
        train = np.squeeze(np.array(train_data[fold]))
        train.sort()
        test = np.squeeze(np.array(test_data[fold]))
        test.sort()

        # generate data splits by year
        ranges, split_train, split_test = split_year_intervals(train, test)

        # generate histograms
        n_subplots = len(split_train)
        fig, axes = plt.subplots(n_subplots, figsize=(30, 15))
        for year in range(n_subplots):
            axes[year].hist(split_train[year], 365, color='g', range=(ranges[year], ranges[year + 1]),
                            label='Train samples')
            axes[year].set_xlim([ranges[year], ranges[year + 1]])
            ax2 = axes[year].twinx()
            ax2.hist(split_test[year], 365, color='r', range=(ranges[year], ranges[year + 1]), label='Test samples')
            axes[year].set_ylabel('Occurences')

        axes[n_subplots - 1].set_xlabel('Days')
        handles_train, labels_train = axes[n_subplots - 1].get_legend_handles_labels()
        handles_test, labels_test = ax2.get_legend_handles_labels()
        axes[n_subplots - 1].legend([handles_train[0], handles_test[0]], [labels_train[0], labels_test[0]])

        plt.suptitle('Distribution of train and test samples', fontsize=28)
        if seed == None:
            fig.savefig(output_path + '/fold_' + str(fold) + '.png')
        else:
            fig.savefig(output_path + '/fold_' + str(fold) + '_%s.png' % str(seed))
        plt.close()

def split_year_intervals(train, test):
    # calculate first and last data sample
    max_date = max(datetime.strptime(train[train.shape[0] - 1], '%y%m%d%H'),
                   datetime.strptime(test[test.shape[0] - 1], '%y%m%d%H'))
    min_date = min(datetime.strptime(train[0], '%y%m%d%H'), datetime.strptime(test[0], '%y%m%d%H'))
    diff_days = (max_date - min_date).days

    # number of years we have data
    n_years = int(diff_days / 365) + 1

    # transform datetimes into days since the first observation
    train_days = []
    for d in train:
        train_days += [(datetime.strptime(d, '%y%m%d%H') - min_date).days]

    test_days = []
    for d in test:
        test_days += [(datetime.strptime(d, '%y%m%d%H') - min_date).days]

    # calculate ranges for yearly plots
    ranges = [x * 365 for x in range(n_years + 1)]
    split_train = [[] for y in range(n_years)]
    split_test = [[] for y in range(n_years)]

    for x in train_days:
        split = int(x / 365)
        split_train[split] += [x]

    for x in test_days:
        split = int(x / 365)
        split_test[split] += [x]

    return ranges, split_train, split_test

# plot the map of switzerland as background image
def plot_map_rlon_rlat(ax, rlat_min, rlat_max, rlon_min, rlon_max, alpha_background):
    dd = xr.open_dataset(
        "/home/n1no/Documents/ethz/master_thesis/code/project/data/c1ffsurf000_timeinvariant_lonlat.nc")
    dd.set_coords(["rlon", "rlat"], inplace=True)
    swiss_data = dd.sel(rlon=slice(rlon_min, rlon_max), rlat=slice(rlat_min, rlat_max))
    swiss_data.FR_LAND.isel(time=0).plot.pcolormesh("rlon", "rlat", ax=ax, alpha=alpha_background,
                                                    cmap=plt.cm.get_cmap('GnBu_r'), add_colorbar=False)
    swiss_data.HH.isel(time=0).plot.pcolormesh("rlon", "rlat", ax=ax, alpha=0.8 * alpha_background,
                                               cmap=plt.cm.get_cmap('YlGn'), add_colorbar=False)

def get_rlon_rlat(grid_points_lon, grid_points_lat):
    dd = xr.open_dataset(
        "/home/n1no/Documents/ethz/master_thesis/code/project/data/c1ffsurf000_timeinvariant_lonlat.nc")
    rlon = dd['rlon'][grid_points_lon]
    rlat = dd['rlat'][grid_points_lat]
    return rlon, rlat

def plot_hour_feature_verification(output_path, hour_feature_dict):
    x, y, color_values = [], [], []

    for key, value in hour_feature_dict.items():
        x.append(key[0])
        y.append(key[1])
        color_values.append(value[0])

    fig = plt.figure(figsize=(15, 15))
    plt.scatter(np.asarray(x), np.asarray(y), c=color_values, s=500, cmap='magma')
    fig.savefig(output_path + '/hour_feature_verification.png')
    plt.close()

def plot_month_feature_verification(output_path, month_feature_dict):
    x, y, color_values = [], [], []

    for key, value in month_feature_dict.items():
        x.append(key[0])
        y.append(key[1])
        color_values.append(value[0])

    fig = plt.figure(figsize=(15, 15))
    plt.scatter(np.asarray(x), np.asarray(y), c=color_values, s=300, cmap='magma')
    fig.savefig(output_path + '/month_feature_verification.png')
    plt.close()

def plotFeatureDistribution(output_path, config,
                            train_features, train_time_invariant_grid_features, train_station_features, train_labels,
                            test_features, test_time_invariant_grid_features, test_station_features, test_labels):
    for idx, train_input in enumerate(train_features):
        fig = plt.figure(figsize=(15, 15))
        plt.hist(train_input, 100)
        plt.suptitle("Train Run %s, Feature %s" % (config['run'], config['features'][idx]))
        fig.savefig(output_path + '/train_run_%s_%s.png' % (config['run'], config['features'][idx]))
        plt.close()

    for idx, train_grid_feature in enumerate(train_time_invariant_grid_features):
        fig = plt.figure(figsize=(15, 15))
        plt.hist(train_grid_feature, 100)
        plt.suptitle(
            "Train Run %s, Time Invariant Grid Feature %s" % (config['run'], config['time_invariant_features'][idx]))
        fig.savefig(output_path + '/train_run_%s_ti_grid_feature_%s.png' % (
            config['run'], config['time_invariant_features'][idx]))
        plt.close()

    for idx, train_station_feature in enumerate(train_station_features):
        fig = plt.figure(figsize=(15, 15))
        plt.hist(train_station_feature, 100)
        plt.suptitle("Train Run %s, StationFeature %s" % (config['run'], config['station_parameters'][idx]))
        fig.savefig(
            output_path + '/train_run_%s_station_feature_%s.png' % (config['run'], config['station_parameters'][idx]))
        plt.close()

    for idx, test_input in enumerate(test_features):
        fig = plt.figure(figsize=(15, 15))
        plt.hist(test_input, 100)
        plt.suptitle("Test Run %s, Feature %s" % (config['run'], config['features']))
        fig.savefig(output_path + '/test_run_%s_%s.png' % (config['run'], config['features'][idx]))
        plt.close()

    for idx, test_grid_feature in enumerate(test_time_invariant_grid_features):
        fig = plt.figure(figsize=(15, 15))
        plt.hist(test_grid_feature, 100)
        plt.suptitle(
            "Test Run %s, Time Invariant Grid Feature %s" % (config['run'], config['time_invariant_features'][idx]))
        fig.savefig(output_path + '/test_run_%s_ti_grid_feature_%s.png' % (
            config['run'], config['time_invariant_features'][idx]))
        plt.close()

    for idx, test_station_feature in enumerate(test_station_features):
        fig = plt.figure(figsize=(15, 15))
        plt.hist(test_station_feature, 100)
        plt.suptitle("Test Run %s, StationFeature %s" % (config['run'], config['station_parameters'][idx]))
        fig.savefig(
            output_path + '/test_run_%s_station_feature_%s.png' % (config['run'], config['station_parameters'][idx]))
        plt.close()

    fig = plt.figure(figsize=(15, 15))
    plt.hist(train_labels, 100)
    fig.savefig(output_path + '/train_labels_run_%s.png' % (config['run']))
    plt.close()

    fig = plt.figure(figsize=(15, 15))
    plt.hist(test_labels, 100)
    fig.savefig(output_path + '/test_labels_run_%s.png' % (config['run']))
    plt.close()

def plotExperimentErrorEvaluation(source_path):
    output_path = source_path + '/plots'
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    error_statistics, epoch_results, best_results, run_times, samples_statistic = [], [], [], [], []
    runs = 0
    for path in glob.glob(source_path + '/**/experiment_statistic.pkl', recursive=True):
        model_name = re.search(r'/([^/]+)/experiment_statistic.pkl', path).group(1)
        with open(path, 'rb') as file:
            ds = pkl.load(file)

        if runs == 0:
            runs = ds.runs

        # get times for which prediction were made
        prediction_times = ds.error_statistic.prediction_times.data

        # get the mean of train and test error over all runs, additionally calculate the std. dev. to the mean
        # run_error_mean_std = tuple(mean, std), mean = ndarray(epoch, error), error = train or test error value,
        # lead = prediction for T = lead
        epoch_error_mean_std = (np.mean(ds.error_statistic.data, 0), np.std(ds.error_statistic.data, 0))
        # look up in train/test mean error per epoch, for which epoch the train resp. test error is the best and
        # return the index
        best_train_epoch, best_test_epoch = np.argmin(epoch_error_mean_std[0], 0)
        # advanced indexing: first select mean or variance. then define fore each prediction lead time the best epoch,
        # the error (train = 0 or test = 1) and the prediction lead time itself.
        # This reduces the [n_epoch, 2, n_future_predictions] -> [n__future_predictions
        best_error_mean_std_epoch = (epoch_error_mean_std[0][best_train_epoch, 0, range(len(prediction_times))],
                                     epoch_error_mean_std[1][best_train_epoch, 0, range(len(prediction_times))],
                                     best_train_epoch,
                                     epoch_error_mean_std[0][best_test_epoch, 1, range(len(prediction_times))],
                                     epoch_error_mean_std[1][best_test_epoch, 1, range(len(prediction_times))],
                                     best_test_epoch)

        # get min, max, mean, std time over all runs and epochs
        time_statistic = ds.run_time_statistic.data
        # first tuple element: name, second tuple element: data [axis 0: [min, max, mean, std], axis 1: epoch, train, test]
        run_times += [(model_name, np.array([np.min(time_statistic, axis=(0, 1)),
                                             np.max(time_statistic, axis=(0, 1)),
                                             np.mean(time_statistic, axis=(0, 1)),
                                             np.std(time_statistic, axis=(0, 1))]))]

        error_statistics += [(model_name, ds.error_statistic.data)]
        epoch_results += [(model_name, epoch_error_mean_std)]
        best_results += [(model_name, best_error_mean_std_epoch)]

        samples_statistic += [(model_name, np.mean(ds.samples_statistic.data, axis=(0, 1)))]

    plotErrorOverEpochsForSingleExperiment(output_path=output_path, error_statistics=error_statistics,
                                           prediction_times=prediction_times)
    plotBestResult(output_path=output_path, best_results=best_results, prediction_times=prediction_times)
    plotCombinedErrorOverEpochs(output_path=output_path, results=epoch_results, prediction_times=prediction_times)
    plotRunTimes(output_path=output_path, timings=run_times)
    plotSampleStatistic(output_path=output_path, samples_statistic=samples_statistic)
    generateTrainingResultTable(output_path=output_path, run_times=run_times, epoch_results=epoch_results, best_results=best_results,
                  prediction_times=prediction_times)

def plotCombinedErrorOverEpochs(output_path, results, prediction_times):
    # plot train errors
    colors = ['r', 'g', 'b', 'y', 'c', 'm', 'y', 'k']
    styles = ['-', '--', ':', '-.']

    plt.rc('lines', linewidth=1)
    plt.rc('axes', prop_cycle=(cycler('color', colors * len(styles)) +
                               cycler('linestyle',
                                      [style for same_style in [[s] * len(colors) for s in styles] for style in
                                       same_style])))

    for prediction_time_idx, prediction_time in enumerate(prediction_times):
        current_output_path = output_path + '/prediction_T%s' % prediction_time
        if not os.path.exists(current_output_path):
            os.makedirs(current_output_path)

        fig, ax = plt.subplots(figsize=(20, 10))

        for name, error in results:
            x = list(range(error[0][:, 0, prediction_time_idx].size))
            ax.plot(x, error[0][:, 0, prediction_time_idx], label=name)

        plt.grid()
        plt.tight_layout(rect=[0, 0.03, 1, 0.95])
        plt.title('Train Error', fontsize=32)
        plt.xticks(x, ['Epoch %s' % i for i in x])
        ax.legend(loc='best', fancybox=True, framealpha=0.4)

        fig.savefig(current_output_path + '/train_error.png')
        plt.close()

        # plot test errors
        plt.rc('lines', linewidth=1)
        plt.rc('axes', prop_cycle=(cycler('color', colors * len(styles)) +
                                   cycler('linestyle',
                                          [style for same_style in [[s] * len(colors) for s in styles] for style in
                                           same_style])))

        fig, ax = plt.subplots(figsize=(20, 10))

        for name, error in results:
            ax.plot(x, error[0][:, 1, prediction_time_idx], label=name)

        plt.grid()
        plt.tight_layout(rect=[0, 0.03, 1, 0.95])
        plt.title('Test Error', fontsize=32)
        plt.xticks(x, ['Epoch %s' % i for i in x])
        ax.legend(loc='best', fancybox=True, framealpha=0.4)

        fig.savefig(current_output_path + '/test_error.png')
        plt.close()

def plotErrorOverEpochsForSingleExperiment(output_path, error_statistics, prediction_times):
    for idx, (name, error_statistic) in enumerate(error_statistics):
        # shape of error statistics: dim_0=Runs, dim_1=Epochs, dim_2=[train_err, test_err]
        runs, epochs, *_ = error_statistic.shape

        for prediction_time_idx, prediction_time in enumerate(prediction_times):

            current_output_path = output_path + '/prediction_T%s' % prediction_time
            if not os.path.exists(current_output_path):
                os.makedirs(current_output_path)

            x = [i + 1 for i in range(epochs)]

            fig, axes = plt.subplots(2, sharey=True, figsize=(20, 20))
            for run in range(runs):
                axes[0].plot(x, error_statistic[run, :, 0, prediction_time_idx], alpha=0.5, label='Run %s' % str(run))
                axes[1].plot(x, error_statistic[run, :, 1, prediction_time_idx], alpha=0.5, label='Run %s' % str(run))
            axes[0].plot(x, np.mean(error_statistic[:, :, 0, prediction_time_idx], axis=0), 'r--', label='Mean')
            axes[1].plot(x, np.mean(error_statistic[:, :, 1, prediction_time_idx], axis=0), 'r--', label='Mean')

            axes[1].set_xlabel('Epoch', fontsize=16)

            axes[0].set_title('Train Error', fontsize=16)
            axes[1].set_title('Test Error', fontsize=16)
            for i in range(2):
                axes[i].set_xticks(x)
                axes[i].set_ylabel('RMSE [°C]', fontsize=16)
                axes[i].grid(True)
                axes[i].legend()

            axes[0].text(0.05, 0.1,
                         'Min. Mean Error: %s' % "{:2.3f}".format(
                             np.min(np.mean(error_statistic[:, :, 0, prediction_time_idx], axis=0))),
                         fontsize=18, fontweight='bold', transform=axes[0].transAxes)
            axes[1].text(0.05, 0.1,
                         'Min. Mean Error: %s' % "{:2.3f}".format(
                             np.min(np.mean(error_statistic[:, :, 1, prediction_time_idx], axis=0))),
                         fontsize=18, fontweight='bold', transform=axes[1].transAxes)
            axes[0].text(0.05, 0.05,
                         'Min. Error: %s' % "{:2.3f}".format(np.min(error_statistic[:, :, 0, prediction_time_idx])),
                         fontsize=18,
                         transform=axes[0].transAxes)
            axes[1].text(0.05, 0.05,
                         'Min. Error: %s' % "{:2.3f}".format(np.min(error_statistic[:, :, 1, prediction_time_idx])),
                         fontsize=18,
                         transform=axes[1].transAxes)

            plt.ylim(0, 1.2 * np.max(error_statistic[:, :, 0, prediction_time_idx]))
            plt.suptitle(str(" ".join(name.split("_")) + " (Lead Prediction = %s)" % prediction_time), fontsize=20)
            fig.tight_layout(rect=[0, 0.03, 1, 0.96], h_pad=4)
            fig.savefig(current_output_path + '/%s_error_plot.png' % name)
            plt.close()

# best results is a list of the best results for each experiment. The best result is a tuple with train in first
# and test in second position. For each of train and test error, the entry is a triple of
# the form (mean error, std. dev., epoch) representing the best result
def plotBestResult(output_path, best_results, prediction_times):
    # create path if it does not exist
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    colors = ['r', 'g', 'b', 'y', 'c', 'm', 'k']

    # plot train error
    plt.rc('lines', linewidth=1)
    plt.rc('axes', prop_cycle=(cycler('color', colors)))

    for prediction_time_idx, prediction_time in enumerate(prediction_times):

        current_output_path = output_path + '/prediction_T%s' % prediction_time
        if not os.path.exists(current_output_path):
            os.makedirs(current_output_path)

        best_train_results = sorted(best_results, key=lambda x: x[1][0][prediction_time_idx])
        fig, ax = plt.subplots(figsize=(20, 10))

        for idx, result in enumerate(best_train_results):
            mean_train = result[1][0]
            std_train = result[1][1]
            mean_test = result[1][3]
            std_test = result[1][4]
            ax.bar(x=idx, height=mean_train, width=0.15, yerr=std_train, label=result[0],
                   color=colors[idx % len(colors)], hatch='x')
            ax.bar(x=idx + 0.2, height=mean_test, width=0.15, yerr=std_test,
                   color=colors[idx % len(colors)], hatch='o')

        ax.set_xticks([])
        plt.grid()
        plt.tight_layout(rect=[0, 0.03, 1, 0.95])
        plt.title('Best Result of all Epochs for each Model averaged over all Runs (sorted by training)', fontsize=32)
        ax.legend(loc='best', fancybox=True, framealpha=0.4)

        fig.savefig(current_output_path + '/best_result_by_train_error.png')
        plt.close()

        # plot test error
        best_test_results = sorted(best_results, key=lambda x: x[1][3][prediction_time_idx])
        fig, ax = plt.subplots(figsize=(20, 10))

        for idx, result in enumerate(best_test_results):
            mean_train = result[1][0]
            std_train = result[1][1]
            mean_test = result[1][3]
            std_test = result[1][4]
            ax.bar(x=idx, height=mean_train, width=0.15, yerr=std_train, label=result[0],
                   color=colors[idx % len(colors)], hatch='x')
            ax.bar(x=idx + 0.2, height=mean_test, width=0.15, yerr=std_test,
                   color=colors[idx % len(colors)], hatch='o')

        ax.set_xticks([])
        plt.grid()
        plt.tight_layout(rect=[0, 0.03, 1, 0.95])
        plt.title('Best Result of all Epochs for each Model averaged over all Runs (sorted by test)', fontsize=32)
        ax.legend(loc='best', fancybox=True, framealpha=0.4)

        fig.savefig(current_output_path + '/best_result_by_test_error.png')
        plt.close()

def plotRunTimes(output_path, timings):
    colors = ['r', 'g', 'b', 'y', 'c', 'm', 'k']
    n_experiments = len(timings)
    x_base = list(range(0, 3))

    fig, ax = plt.subplots(figsize=(20, 10))
    for idx, (model_name, timing) in enumerate(timings):
        shift = (idx / n_experiments / 2)
        x = [pos + shift for pos in x_base]
        ax.bar(x=x, height=timing[0:3, 0], width=0.05, label=model_name, color=colors[idx % len(colors)])

    plt.xticks([pos + ((n_experiments - 1) / 2) / n_experiments / 2 for pos in x_base],
               ('Fastest Time', 'Slowest Time', 'Mean Time'))
    plt.grid()
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    plt.title('Min/Max/Mean Time over all Epochs and Runs', fontsize=32)
    ax.legend(loc='best', fancybox=True, framealpha=0.4)

    fig.savefig(output_path + '/epoch_run_times.png')
    plt.close()

def plotSampleStatistic(output_path, samples_statistic):
    # ------------------------------------------------------------------------------------------------------------------------
    # skipped samples statistics due to corrupted measurement stations
    # ------------------------------------------------------------------------------------------------------------------------

    fig = plt.figure(figsize=(20, 10))
    N = len(samples_statistic)
    ind = np.arange(N)  # the x locations for the groups
    experiment_title = [splitModelNameVertical(name) for name, _ in samples_statistic]
    data = np.zeros((4, N))
    for idx, (_, samples) in enumerate(samples_statistic):
        data[:, idx] = np.array([samples[1], samples[0] - samples[1], samples[3], samples[2] - samples[3]])

    width = 0.35  # the width of the bars: can also be len(x) sequence

    p1 = plt.bar(ind, data[0, :], width)
    p2 = plt.bar(ind, data[1, :], width, bottom=data[0, :])
    p3 = plt.bar(ind, data[2, :], width, bottom=np.sum(data[0:2, :], axis=0))
    p4 = plt.bar(ind, data[3, :], width, bottom=np.sum(data[0:3, :], axis=0))

    plt.ylabel('Samples')
    plt.title('Processed / Skipped Samples', fontsize=22)
    plt.xticks(ind, experiment_title)
    plt.xlim((-0.5, np.max(ind) + 1))
    plt.legend((p1[0], p2[0], p3[0], p4[0]), ('Processed Train', 'Skipped Train', 'Processed Test', 'Skipped Test'),
               loc='upper right')
    fig.savefig(output_path + '/skip_sample_statistic.png')
    plt.close()

def splitModelNameVertical(name):
    n = 5
    split = name.split('_')
    lines = []
    for i in range(0, len(split), n):
        lines += [" ".join(split[i:i + n])]
    return "\n".join(lines)

def generateTrainingResultTable(output_path, run_times, epoch_results, best_results, prediction_times):
    with open(output_path + '/table_data.csv', "w") as csv_file:
        title_per_epoch = ["Train Mean, Train Std., Test Mean, Test Std."]
        csv_file.write(",".join(["Name, Lead Prediction"] + title_per_epoch * epoch_results[0][1][0].shape[0] + [
            "Best Train Mean, Beat Train Std., Best Test Mean, Best Test Std.\n"]))

        for prediction_time_idx, prediction_time in enumerate(prediction_times):
            for idx, (name, epoch_result) in enumerate(epoch_results):
                # first stack train mean error and std to one tuple and test mean error and std to other tuple, resulting
                # in array([[mean_train_ep_1, std_train_ep_1], ..., [mean_train_ep_n, mean_train_ep_n]]) and
                # array([[mean_test_ep_1, std_test_ep_1], ..., [mean_test_ep_n, mean_test_ep_n]]). Stack those two
                # horizontally to get array([[mean_train_ep_1, std_train_ep_1, mean_test_ep_1, std_test_ep_1], ...,
                # [mean_train_ep_n, mean_train_ep_n, mean_test_ep_n, mean_test_ep_n]]). Finally, flatten to list.
                results = [str(elem) for elem in
                           np.hstack((np.stack(epoch_result, axis=3)[:, 0, prediction_time_idx], np.stack(epoch_result, axis=3)[:, 1, prediction_time_idx])).flatten().tolist()]
                best_train_mean, best_train_std, _, best_test_mean, best_test_std, _ = tuple(map(lambda arr: arr[prediction_time_idx], best_results[idx][1]))

                csv_file.write(",".join([name, str(prediction_time)] + results + [str(elem) for elem in
                                                                                [best_train_mean, best_train_std,
                                                                                best_test_mean,
                                                                                best_test_std]] + ["\n"]))

# plot the error from the bias corrected baseline run. The error statistic has the following dimensions (runs, errors),
# where errors = [raw_error, station_bias_corrected, station_hour_bias_corrected, station_hour_month_bias_corrected']
def plotBiasCorrectedBaselineResults(source_path):
    output_path = source_path + '/plots'
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    error_data = {
        '2d': None,
        '3d': None
    }
    for path in glob.glob(source_path + '/**/experiment_statistic.pkl', recursive=True):
        with open(path, 'rb') as file:
            ds = pkl.load(file)
        distance_metric = ds.attrs['experiment_info']['distance_metric']
        error_statistic = ds.error_statistic

        try:
            error_data[distance_metric] = xr.concat([error_data[distance_metric], error_statistic], dim='run')
        except TypeError:
            error_data[distance_metric] = error_statistic

    lead_error = {}
    for metric_idx, distance_metric_key in enumerate(error_data.keys()):
        error_by_metric = error_data[distance_metric_key]
        if error_by_metric is None: continue
        mean_squared_error_data = np.sqrt(np.nanmean(np.square(error_by_metric.data), axis=(2, 3, 4, 6)))
        lead_error[distance_metric_key] = mean_squared_error_data

        # TODO nwi plot stuff...

def plotPredictionRun(source_path, observation_path, n_parallel=1):
    plotAveragedPredictionRun(source_path)
    plotPerStationPredictionRun(source_path, observation_path, n_parallel)

def plotAveragedPredictionRun(source_path):
    # gather all models in source folder
    error_data_per_run_dict = {}
    for path in glob.glob(source_path + '/**/model_run_error.pkl', recursive=True):
        model_name = path.split('/')[-2]
        with open(path, 'rb') as file:
            ds = pkl.load(file)
        for data_var in ds.data_vars:
            inits = ds[data_var].init.data
            sample_type_mapping = [mapping[1] for mapping in ds[data_var].init_type_mapping]
            prediction_data = ds[data_var].data
            try:
                error_data_per_run_dict[data_var] += [(model_name, inits, prediction_data, sample_type_mapping)]
            except:
                error_data_per_run_dict[data_var] = [(model_name, inits, prediction_data, sample_type_mapping)]

    # get the prediction lead time to adjust time labels
    prediciton_lead_time = ds.attrs['config']['prediction_times'][0] if 'config' in ds.attrs else 1

    times = DataUtils.getTimeFromFileName(inits, prediciton_lead_time)
    time_labels = [str(t)[:-13] for t in times]

    for run_error_data in error_data_per_run_dict.items():
        run = run_error_data[0]
        model_mean_errors = {}
        n_subplots = 10
        fig, axes = plt.subplots(n_subplots, figsize=(60, 20), sharey=True)
        for model_idx, model_error_data in enumerate(run_error_data[1]):

            N = len(model_error_data[1])
            split_length = N // n_subplots
            ind = np.arange(N)  # the x locations for the groups
            experiment_title = model_error_data[0]
            prediction_data = model_error_data[2]
            init_type_mapping = model_error_data[3]
            train_indices = [idx for idx, item in enumerate(init_type_mapping) if item == 'train']
            test_indices = [idx for idx, item in enumerate(init_type_mapping) if item == 'test']
            filtered_indices = [idx for idx, item in enumerate(init_type_mapping) if item == 'filterd']

            for i in range(n_subplots):
                # split indexes into slices for each subplot
                index_split = ind[i * split_length:(i + 1) * split_length]

                if model_idx == 0:
                    sampleTypeBackgroundColoring(axes[i], index_split,
                                                 init_type_mapping[i * split_length:(i + 1) * split_length])
                    axes[i].set_xlim([np.min(index_split), np.max(index_split)])

                axes[i].plot(index_split,
                             np.nanmean(prediction_data[i * split_length:(i + 1) * split_length,:, 0],axis=1),
                             label=experiment_title, linewidth=0.15, alpha=0.8)

            train_model_bias = np.nanmean(prediction_data[train_indices][:,:,3])
            train_model_rmse = np.sqrt(np.nanmean(np.square(prediction_data[train_indices][:,:,3])))
            train_model_mae = np.nanmean(np.absolute(prediction_data[train_indices][:,:,3]))

            test_model_bias = np.nanmean(prediction_data[test_indices][:,:,3])
            test_model_rmse = np.sqrt(np.nanmean(np.square(prediction_data[test_indices][:,:,3])))
            test_model_mae = np.nanmean(np.absolute(prediction_data[test_indices][:,:,3]))

            filtered_model_bias = np.nanmean(prediction_data[filtered_indices][:,:,3])
            filtered_model_rmse = np.sqrt(np.nanmean(np.square(prediction_data[filtered_indices][:,:,3])))
            filtered_model_mae = np.nanmean(np.absolute(prediction_data[filtered_indices][:,:,3]))

            model_mean_errors[experiment_title] = (train_model_bias, train_model_rmse, train_model_mae,
                                                   test_model_bias, test_model_rmse, test_model_mae,
                                                   filtered_model_bias, filtered_model_rmse, filtered_model_mae)

        # add mean errors of cosmo output predictions
        train_diff_cosmo = prediction_data[train_indices][:,:,1] - prediction_data[train_indices][:,:,2]
        train_cosmo_bias = np.nanmean(train_diff_cosmo)
        train_cosmo_rmse = np.sqrt(np.nanmean(np.square(train_diff_cosmo)))
        train_cosmo_mae = np.nanmean(np.absolute(train_diff_cosmo))

        test_diff_cosmo = prediction_data[test_indices][:,:,1] - prediction_data[test_indices][:,:,2]
        test_cosmo_bias = np.nanmean(test_diff_cosmo)
        test_cosmo_rmse = np.sqrt(np.nanmean(np.square(test_diff_cosmo)))
        test_cosmo_mae = np.nanmean(np.absolute(test_diff_cosmo))

        filtered_diff_cosmo = prediction_data[filtered_indices][:,:,1] - prediction_data[filtered_indices][:,:,2]
        filtered_cosmo_bias = np.nanmean(filtered_diff_cosmo)
        filtered_cosmo_rmse = np.sqrt(np.nanmean(np.square(filtered_diff_cosmo)))
        filtered_cosmo_mae = np.nanmean(np.absolute(filtered_diff_cosmo))

        # add COSMO-1 output prediction error
        model_mean_errors['COSMO-1'] = (train_cosmo_bias, train_cosmo_rmse, train_cosmo_mae,
                                        test_cosmo_bias, test_cosmo_rmse, test_cosmo_mae,
                                        filtered_cosmo_bias, filtered_cosmo_rmse, filtered_cosmo_mae)

        for i in range(n_subplots):
            axes[i].plot(ind[i * split_length:(i + 1) * split_length],
                         np.nanmean(prediction_data[i * split_length:(i + 1) * split_length,:, 1], axis=1), label='COSMO-1',
                         linewidth=0.15, alpha=0.8, color='b', linestyle='-.')
            axes[i].plot(ind[i * split_length:(i + 1) * split_length],
                         np.nanmean(prediction_data[i * split_length:(i + 1) * split_length, 2], axis=1), label='Prediction',
                         linewidth=0.15, alpha=0.8, color='m', linestyle='--')

            tick_step_size = np.maximum(split_length // 30, 1)
            axes[i].set_xticks(ind[i * split_length:(i + 1) * split_length][::tick_step_size])
            axes[i].set_xticklabels(time_labels[i * split_length:(i + 1) * split_length][::tick_step_size])
            axes[i].set_xticks(ind[i * split_length:(i + 1) * split_length], minor=True)
            # And a corresponding grid
            axes[i].grid(which='both')

            # Or if you want different settings for the grids:
            axes[i].grid(which='minor', alpha=0.2)
            axes[i].grid(which='major', alpha=0.5)

            handles, labels = axes[0].get_legend_handles_labels()

        axes[n_subplots - 1].set_xlabel('Time')
        axes[0].legend(handles, labels)
        plt.tight_layout()

        run_path = source_path + '/plots/prediction_runs/%s' % run
        if not os.path.exists(run_path):
            os.makedirs(run_path)

        fig.savefig(run_path + '/averaged_prediction.png', dpi=300)

        generatePredictionResultTable(output_path=run_path, results=model_mean_errors)

def plotPerStationPredictionRun(source_path, observation_path, n_parallel):
    # gather all models in source folder
    error_data_per_run_dict = defaultdict()
    for path in glob.glob(source_path + '/**/model_run_error.pkl', recursive=True):
        model_name = path.split('/')[-2]
        with open(path, 'rb') as file:
            ds = pkl.load(file)
        for data_var in ds.data_vars:
            da = ds[data_var]
            try:
                error_data_per_run_dict[data_var] += [(model_name, da)]
            except:
                error_data_per_run_dict[data_var] = [(model_name, da)]

    # load observations
    OBS = xr.open_dataset(observation_path)
    # get the prediction lead time to adjust time labels
    prediciton_lead_time = ds.attrs['config']['prediction_times'][0] if 'config' in ds.attrs else 1

    for run_error_data in error_data_per_run_dict.items():
        run = run_error_data[0]
        models = run_error_data[1]
        stations = run_error_data[1][0][1].station.data
        inits = run_error_data[1][0][1].init.data
        init_type_mapping = np.array(run_error_data[1][0][1].init_type_mapping)
        train_indices = [idx for idx, item in enumerate(init_type_mapping) if item[1] == 'train']
        test_indices = [idx for idx, item in enumerate(init_type_mapping) if item[1] == 'test']
        sample_type_color_mapping = [mapping[1] for mapping in init_type_mapping]
        times = DataUtils.getTimeFromFileName(inits, prediciton_lead_time)
        time_labels = [str(t)[:-13] for t in times]

        station_name_dict = get_station_dict(OBS, stations)

        model_station_mean_errors = {}
        # plot for each station the prediction run results in parallel
        with Pool(processes=n_parallel) as pool:
            process_results = []

            for station_idx, station in enumerate(stations):
                print('Plotting of prediction run for station %s queued.' % station)
                process_results.append(pool.apply_async(plotPerStationPredictionRunWorker,
                                                        (models, station, train_indices, test_indices,
                                                         station_name_dict,sample_type_color_mapping,
                                                         time_labels, source_path, run)))

            # aggregate results from all processes
            for ps_idx, ps_result in enumerate(process_results):
                # sync processes
                model_station_mean_error = ps_result.get()

                for experiment_title, station_data_list in model_station_mean_error.items():
                    try:
                        model_station_mean_errors[experiment_title] += station_data_list
                    except KeyError:
                        model_station_mean_errors[experiment_title] = station_data_list

                print('[Process %s] Synchronized after plotting station.' % ps_idx)

        run_path = source_path + '/plots/prediction_runs/%s' % run
        if not os.path.exists(run_path):
            os.makedirs(run_path)

        generateStationPredictionResultTable(output_path=run_path, results=model_station_mean_errors)

def plotPerStationPredictionRunWorker(models, station, train_indices, test_indices, station_name_dict,
                                      sample_type_color_mapping, time_labels, source_path, run):
    n_subplots = 10
    model_station_mean_error = {}
    fig, axes = plt.subplots(n_subplots, figsize=(60, 20), sharey=True)

    for model_idx, model_error_data in enumerate(models):
        experiment_title = model_error_data[0]

        station_error_data = model_error_data[1].sel(station=station).data

        train_station_bias = np.nanmean(station_error_data[train_indices][:, 3])
        train_station_rmse = np.sqrt(np.nanmean(np.square(station_error_data[train_indices][:, 3])))
        train_station_mae = np.nanmean(np.absolute(station_error_data[train_indices][:, 3]))

        test_station_bias = np.nanmean(station_error_data[test_indices][:, 3])
        test_station_rmse = np.sqrt(np.nanmean(np.square(station_error_data[test_indices][:, 3])))
        test_station_mae = np.nanmean(np.absolute(station_error_data[test_indices][:, 3]))

        model_station_mean_error[experiment_title] = [(station_name_dict[station], station,
                                                        train_station_bias, train_station_rmse,
                                                        train_station_mae, test_station_bias,
                                                        test_station_rmse, test_station_mae)]

        N = len(station_error_data)
        split_length = N // n_subplots
        ind = np.arange(N)  # the x locations for the groups
        for i in range(n_subplots):
            # split indexes into slices for each subplot
            index_split = ind[i * split_length:(i + 1) * split_length]

            # for the first model we plot the background coloring
            # denoting train, test and filtered data sections
            if model_idx == 0:
                sampleTypeBackgroundColoring(axes[i], index_split,
                                             sample_type_color_mapping[i * split_length:(i + 1) * split_length])
                axes[i].set_xlim([np.min(index_split), np.max(index_split)])

            # plot network outpu
            axes[i].plot(index_split,
                         station_error_data[i * split_length:(i + 1) * split_length, 0],
                         label=experiment_title, linewidth=0.15, alpha=0.8)

    for i in range(n_subplots):
        # plot cosmo prediction
        axes[i].plot(ind[i * split_length:(i + 1) * split_length],
                     station_error_data[i * split_length:(i + 1) * split_length, 1],
                     label='COSMO-1', linewidth=0.15, alpha=0.8, color='b', linestyle='-.')

        axes[i].plot(ind[i * split_length:(i + 1) * split_length],
                     station_error_data[i * split_length:(i + 1) * split_length, 2],
                     label='Target', linewidth=0.15, alpha=0.8, color='m', linestyle='--')

        tick_step_size = np.maximum(split_length // 30, 1)
        axes[i].set_xticks(ind[i * split_length:(i + 1) * split_length][::tick_step_size])
        axes[i].set_xticklabels(time_labels[i * split_length:(i + 1) * split_length][::tick_step_size])
        axes[i].set_xticks(ind[i * split_length:(i + 1) * split_length], minor=True)
        # And a corresponding grid
        axes[i].grid(which='both')

        # Or if you want different settings for the grids:
        axes[i].grid(which='minor', alpha=0.2)
        axes[i].grid(which='major', alpha=0.5)

        handles, labels = axes[0].get_legend_handles_labels()

    # add error of cosmo prediction
    train_diff_cosmo = station_error_data[train_indices][:, 1] - station_error_data[train_indices][:, 2]
    train_station_bias_cosmo = np.mean(train_diff_cosmo)
    train_station_rmse_cosmo = np.sqrt(np.nanmean(np.square(train_diff_cosmo)))
    train_station_mae_cosmo = np.nanmean(np.absolute(train_diff_cosmo))

    test_diff_cosmo = station_error_data[test_indices][:, 1] - station_error_data[test_indices][:, 2]
    test_station_bias_cosmo = np.nanmean(test_diff_cosmo)
    test_station_rmse_cosmo = np.sqrt(np.nanmean(np.square(test_diff_cosmo)))
    test_station_mae_cosmo = np.nanmean(np.absolute(test_diff_cosmo))

    model_station_mean_error['COSMO-1'] = [(station_name_dict[station], station,
                                             train_station_bias_cosmo, train_station_rmse_cosmo,
                                             train_station_mae_cosmo, test_station_bias_cosmo,
                                             test_station_rmse_cosmo, test_station_mae_cosmo)]

    axes[n_subplots - 1].set_xlabel('Time')
    axes[0].legend(handles, labels)
    plt.tight_layout()

    run_path = source_path + '/plots/prediction_runs/%s' % run
    if not os.path.exists(run_path):
        os.makedirs(run_path)

    fig.savefig(run_path + '/prediction_station_%s_%s' % (
    station, station_name_dict[station].replace("/", "").replace("  ", " ")) + '.png', dpi=300)

    return model_station_mean_error

def generatePredictionResultTable(output_path, results):
    with open(output_path + '/prediction_results_table_data.csv', "w") as csv_file:
        # write csv title
        csv_file.write("Experiment, Train Station Bias, Train Station RMSE, Train Station MAE, "
                       "Test Station Bias, Test Station RMSE, Test Station MAE,"
                       "Filtered Station Bias, Filtered Station RMSE, Filtered Station MAE,\n")

        for model_name, model_result in results.items():
                csv_file.write(",".join([model_name] + [str(result) for result in model_result] + ["\n"]))

def generateStationPredictionResultTable(output_path, results):
    with open(output_path + '/prediction_station_results_table_data.csv', "w") as csv_file:
        # write csv title
        csv_file.write("Experiment, Station Name, Station Id, Train Station Bias, Train Station RMSE, Train Station MAE, Test Station Bias, Test Station RMSE, Test Station MAE, \n")

        for model_name, model_result in results.items():
            for station_result_idx, station_result in enumerate(model_result):
                csv_file.write(",".join([model_name] + [str(result) for result in station_result] + ["\n"]))

def sampleTypeBackgroundColoring(ax, index, sample_type):
    assert len(index) == len(sample_type)
    prev_type = None
    coloring_edge_points = []
    for idx, type in enumerate(sample_type):
        if prev_type != type:
            coloring_edge_points += [(index[idx], type)]
            prev_type = type
    coloring_edge_points += [(index[-1], None)]

    last_point = coloring_edge_points[0]
    for edge_point_index in range(1, len(coloring_edge_points)):
        next_point = coloring_edge_points[edge_point_index]
        ax.axvspan(last_point[0], next_point[0], facecolor=type_color_mapping[last_point[1]], alpha=0.05)
        last_point = next_point

if __name__ == '__main__':
    experiment = sys.argv[1]
    if experiment == "network":
        plotExperimentErrorEvaluation(source_path=sys.argv[2])
    elif experiment == "baseline":
        plotBiasCorrectedBaselineResults(source_path=sys.argv[2])
    elif experiment == "plotPredictionRun":
        plotPredictionRun(source_path=sys.argv[2], observation_path=sys.argv[3])


import logging
from pythonjsonlogger import jsonlogger

def log_setup(level = logging.INFO):
    root_logger= logging.getLogger()
    root_logger.setLevel(level)
    handler = logging.FileHandler('error.log', 'a', 'utf-8')
    handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(name)s %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p'))
    root_logger.addHandler(handler)

def json_log_setup(level = logging.DEBUG):
    logger = logging.getLogger()
    if logger.hasHandlers():
        logger.handlers.clear()
    logger.setLevel(level)
    logHandler = logging.FileHandler('error_json.log', 'a', 'utf-8')
    formatter = jsonlogger.JsonFormatter('%(asctime)s %(levelname)s %(name)s %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p')
    logHandler.setFormatter(formatter)
    logger.addHandler(logHandler)


#!/usr/bin/env python3

if 1:
    # This is needed for the SE calculation
    import rpy2
    from rpy2.robjects.packages import importr
    from rpy2.robjects import FloatVector
    from rpy2.robjects import r
    importr('coda')

import statistics
import math
import numpy
# from scipy.stats import beta
import sys
import argparse

# https://docs.python.org/3/howto/argparse.html
# https://docs.python.org/3/library/argparse.html
# https://docs.python.org/3/library/argparse.html#the-add-argument-method

myDescription = """Digests log likelihoods for stepping stone method.

Wangang Xie, Paul O. Lewis, Yu Fan, Lynn Kuo, Ming-Hui Chen 2011.
Improving Marginal Likelihood Estimation for Bayesian Phylogenetic Model Selection.
Systematic Biology, 60 150-–160 https://doi.org/10.1093/sysbio/syq085

Thanks to Mario dos Reis for explanations and code on his website, https://dosreislab.github.io/2018/08/31/bppr-stepstones.html

For the variance calculation, this implementation uses the effectiveSize function in the coda package in R, so that needs to be available on your system, as well as the rpy2 Python module (https://rpy2.github.io/doc.html) to enable Python to use R.

The log likelihoods need to be all in a single column in one file, and should go from high (all Posterior) to low beta (all Prior) values.

As in Xie et al, the beta values are given as β_k = (k/K)^{1/α}, where alpha is less than one.  That is a different parameterization to the Beta(alpha,1.0) as given in Scipy.

If you provide an alpha value, this script will use the Xie et al β_k formula to to calculate beta values.  Alternatively, if your beta values are in a file, provide the file name.  The latter option also provides K.

"""

# See https://stackoverflow.com/questions/35857610/how-to-preserve-newlines-in-argparse-version-output-while-letting-argparse-auto
# for info about RawDescriptionHelpFormatter
parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,description=myDescription)

parser.add_argument("logLikesFile", type=str, help="Name of the file containing the log likelihoods")
parser.add_argument("-sps", "--samplesPerStep", type=int, help="The number of power posteriors per step")
parser.add_argument("-K", "--K", type=int, help="K, the number of stepping stones ")
parser.add_argument("-alpha", "--alpha", type=float, help="The alpha for beta_k values, as in Xie et al pg 154, eg 0.25 or 0.3")
parser.add_argument("-bt", "--betaValsFile", type=str, help="Name of the file containing the beta values")
parser.add_argument("-doSE", "--doSE", action="store_true", help="Calculate the SE (standard error).  Optional --- default is False")
parser.add_argument("-gr", "--gram", action="store_true", help="Plot using gram, default is False. Gram.baseName is meanLogLikes")
# parser.add_argument("-v", "--verbosity", type=int, choices=[0, 1, 2], default=0, help="requires an argument, must be one of the choices")

# If there is no point in proceeding ...
if len(sys.argv)==1:
    parser.print_help()
    sys.exit(1)

args = parser.parse_args()
# print(args)

assert args.samplesPerStep, "Need to specify the --samplesPerStep, the number of samples in each stepping stone"

# Get the beta values, either from a file, or regenerate them.
if args.betaValsFile:
    bb = [float(l) for l in open("beta.txt")]
    K = len(bb)
else:
    assert args.alpha
    assert args.K

    alpha = args.alpha
    K = args.K
    bb = [math.pow((k/K),(1/alpha)) for k in range(K)]
    #print("bb ", bb)

bb.append(1.0)
bdiffs = [bb[i] - bb[i-1] for i in range(1, K + 1)]

maxima = numpy.zeros(K)
zr = numpy.zeros(K)           # mean scaled likelihoods (not log likelihoods)
meanLogLikes = numpy.zeros(K)
ess = numpy.zeros(K)
vzr = numpy.zeros(K)
likes = []

#print(f"bb {bb}")
#print(f"bdiffs {bdiffs}")

# All likelihoods from all beta values
allLikes = [float(l) for l in open(args.logLikesFile)]
myMsg = f"Got {len(allLikes)} log likelihood values, should be samplesPerStep {args.samplesPerStep} * K {K} = {args.samplesPerStep * K}"
assert len(allLikes) == args.samplesPerStep * K, myMsg
allLikes.reverse()

for k in range(K):
    # Get the log likes into a numpy array, logLikes
    start = args.samplesPerStep * k
    logLikes = numpy.array(allLikes[start: start + args.samplesPerStep])

    meanLogLikes[k] = statistics.mean(logLikes)
    #print(f"meanLogLikes {meanLogLikes}")

    logLikes *= bdiffs[k]
    maxima[k] = max(logLikes)
    #print(f"maxima {maxima}")

    likes.append(numpy.exp(logLikes - maxima[k]))

    zr[k] = numpy.mean(likes[k])

for k in range(K):
    # From Mario
    #ess[k] <- coda::effectiveSize(likes[[k]])
    fvR = FloatVector(likes[k])
    ret = r['effectiveSize'](fvR)
    myFloat = float(ret[0])
    ess[k] = myFloat

    vzr[k] = numpy.var(likes[k]) / ess[k]

    # the delta approximation does not work well if vzr/zr^2 > 0.1
    if (vzr[k] / (zr[k] * zr[k])) > 0.1:
       print("unreliable se: var(r_k)/r_k^2 = ", vzr[k] / (zr[k] * zr[k]), f" > 0.1 for k={k}, b=", bb[k])

if 0:
    print("%3s %12s %12s %12s %12s %12s %14s" % ("k", "b", "var(likes)", "ess", "vzr", "zr", "var(r_k)/r_k^2"))
    for k in range(K):
        print(f"{k:3}", end=' ')
        print(f"{bb[k]:12.6}", end=' ')
        print(f"{numpy.var(likes[k]):12.4f}", end=' ')
        print(f"{ess[k]:12.4f}", end=' ')
        print(f"{vzr[k]:12.8f}", end=' ')
        print(f"{zr[k]:12.6f}", end=' ')
        print(f"{vzr[k] / (zr[k] * zr[k]):14.6f}", end=' ')
        print()

vmlnl = sum(vzr / (zr * zr))

lnml = sum(numpy.log(zr) + maxima)

if args.gram:
    from gram import Plot
    gp = Plot()
    xx = list(bb[:-1])
    yy = list(meanLogLikes)
    #gp.svgPxForCm = 100
    gp.baseName = 'meanLogLikes'
    #gp.scatter(xx, yy)
    gp.line(xx, yy)
    gp.xAxis.title = r'$\beta_k$'
    gp.yAxis.title = 'mean log likelihood'
    gp.minXToShow = 0
    gp.maxXToShow = 1
    gp.pdf()

#print(meanLogLikes)
#print(maxima)
#print(zr)
#print(ess)
#print(vzr)
# print(vmlnl)
print("log marginal likelihood:", lnml)


import tensorflow as tf
from tensorflow.keras import layers, optimizers
import nltk
import refine_data_set
import data_organise
import text_vectorise
import random

lemmatizer = nltk.stem.WordNetLemmatizer()
stop_words = nltk.corpus.stopwords.words('english')

if __name__ == '__main__':
    raw_X, raw_Y = refine_data_set.refine_X_Y()

    balanced_X, balanced_Y = data_organise.balance_data_binary(raw_X, raw_Y)

    # count = 0
    # for i in range(10000):
    #     j = random.randint(0, len(balanced_X) - 1)
    #     if raw_Y[raw_X.index(balanced_X[j])] != balanced_Y[j]:
    #         print(raw_Y[raw_X.index(balanced_X[j])])
    #         print(balanced_Y[j])
    #         count += 1
    # print('Missed ', count, ' in 10000')

    # print('There are ', balanced_Y.count(0), '\'0s\' in balanced_Y')
    # print('There are ', balanced_Y.count(1), '\'1s\' in balanced_Y')

    X_vectorised, word_index_map = text_vectorise.vectorise(balanced_X, lemmatizer, stop_words, 2, twit=True)
    # split_train_test shuffles and splits

    X_train, Y_train, X_test, Y_test = data_organise.split_train_test(X_vectorised, balanced_Y, split=0.7)

    print("number of training examples = " + str(X_train.shape[1]))
    print("number of test examples = " + str(X_test.shape[1]))
    print("X_train shape: " + str(X_train.shape) + ' of type: ', type(X_train))
    print("Y_train shape: " + str(Y_train.shape) + ' of type: ', type(Y_train))
    print("X_test shape: " + str(X_test.shape) + ' of type: ', type(X_test))
    print("Y_test shape: " + str(Y_test.shape) + ' of type: ', type(Y_test))

    model = tf.keras.Sequential()
    model.add(layers.Dense(25, activation='relu', input_dim=X_vectorised.shape[1]))
    model.add(layers.Dense(12, activation='relu'))
    model.add(layers.Dense(1, activation='sigmoid'))

    model.summary()

    model.compile(optimizer=optimizers.Adam(learning_rate=0.001), loss='binary_crossentropy', metrics=['accuracy'])

    model.fit(X_train, Y_train, batch_size=64, epochs=50, validation_data=(X_test, Y_test))


from rest_framework import permissions
from django.contrib.auth.models import Permission
from polls.models import Course, UserRole

# from django.shortcuts import get_object_or_404

# class IsInstructor(permissions.IsAuthenticated):
#     """
#     permission check for if user is instructor
#     """

#     def has_permission(self, request, view):
#         return super().has_permission(request, view) and request.user.profile.is_instructor

class IsInstructorOrAdmin(permissions.IsAuthenticated):
    """
    permission check for if user is instructor or admin.
    """

    def has_permission(self, request, view):
        # print(request.query_params)
        if request.user.is_staff:
            return True
        pk = request.data.get('course', None)
        if pk is None:
            pk = request.data.get('pk', None)
        if pk is None:
            pk = request.data.get('id', None)
        if pk is None:
            pk = dict(request.query_params).get('courses[]', None)
            if pk is not None:
                pk = int(pk[0])
        print("course pk={}".format(pk))
        if super().has_permission(request, view) and pk is not None:
            course = Course.objects.get(pk=pk)
            try:
                role = UserRole.get(user=request.user, course=course).role
                perm = Permission.get(codename='view_question')
                print(role.permissions.all())
                print(perm)
                if perm in role.permissions.all():
                    return True
            except UserRole.DoesNotExist:
                pass
        return False


from collections import deque

from search_matcher import SearchMatcher

# Number of random strings to bootstrap the search
start_count  = 100
# Length of start strings
start_length = 5
# The string we are trying to find
target = "justin"

# List of SearchMatchers.  The work to do.
Q = deque()

# Create the initial strings
for i in range(0, start_count):
    Q.append(SearchMatcher(target, None, None, start_length))
print(Q)

# Max number of strings to evaluate before giving up
max_tries = 1*1000*1000
# Number of tries so far
tries = 0

# Did we find the target?
success = False
while True:

    print(len(Q))
    try:
        current = Q.popleft()
    except IndexError:
        break
    if tries == max_tries:
        break
    # new: List of new SearchMatchers
    new = current.decide()
    if new == SearchMatcher.FOUND:
        success = True
        break
    # Add new candidate strings:
    Q.extend(new)
    tries += 1

# print(Q)
# Dump the remaining Q in order of score:
print("queue:")
L = sorted(Q)
for item in L:
    print(item)
print("length: %i" % len(Q))
print("result: " + str(success))


import argparse
import atexit
import logging
import math
import os
import sys
import time
import numpy as np

from numpy import genfromtxt
import scipy.signal as signal

from cyber.python.cyber_py3 import cyber
from cyber.python.cyber_py3 import cyber_time
from modules.tools.common.logger import Logger
from modules.canbus.proto import chassis_pb2
from modules.common.configs.proto import vehicle_config_pb2
from modules.common.proto import drive_state_pb2
from modules.common.proto import pnc_point_pb2
from modules.control.proto import pad_msg_pb2
from modules.localization.proto import localization_pb2
from modules.planning.proto import planning_pb2
from modules.prediction.proto import  prediction_obstacle_pb2
import modules.tools.common.proto_utils as proto_utils
from modules.dreamview.proto import chart_pb2

## local package
from Agent.zzz.dynamic_map import DynamicMap
from Agent.zzz.frenet import Frenet_path
from Agent.zzz.JunctionTrajectoryPlanner import JunctionTrajectoryPlanner
from Agent.zzz.JunctionTrajectoryPlanner_simple_predict import JunctionTrajectoryPlanner_SP
from Agent.zzz.JunctionTrajectoryPlanner_simple_predict_dns import JunctionTrajectoryPlanner_SP_Dns
from Agent.zzz.dynamic_map import Lane, Lanepoint, Vehicle
from scipy.spatial.transform import Rotation as R
from modules.routing.proto import routing_pb2
from modules.perception.proto import perception_obstacle_pb2
from matplotlib import pyplot as plt

# TODO(all): hard-coded path temporarily. Better approach needed.
APOLLO_ROOT = "/apollo"

class RtkPlayer(object):
    """
    rtk player class
    """

    def __init__(self, record_file, node):
        """Init player."""
        self.firstvalid = False
        self.logger = Logger.get_logger(tag="RtkPlayer")
        self.logger.info("Load record file from: %s" % record_file)
        try:
            file_handler = open(record_file, 'r')
        except (FileNotFoundError, IOError) as ex:
            self.logger.error("Error opening {}: {}".format(record_file, ex))
            sys.exit(1)

        self.data = genfromtxt(file_handler, delimiter=',', names=True)
        file_handler.close()
        self.starttime = cyber_time.Time.now().to_sec()
        self.localization = localization_pb2.LocalizationEstimate()
        self.prediction = prediction_obstacle_pb2.PredictionObstacles()
        self.perception = perception_obstacle_pb2.PerceptionObstacles()
        self.chassis = chassis_pb2.Chassis()
        self.padmsg = pad_msg_pb2.PadMessage()
        self.localization_received = False
        self.chassis_received = False

        self.planning_pub = node.create_writer('/apollo/planning',
                                               planning_pb2.ADCTrajectory)

        self.speedmultiplier = 1
        self.terminating = False
        self.sequence_num = 0

        b, a = signal.butter(6, 0.05, 'low')
        self.data['acceleration'] = signal.filtfilt(b, a,
                                                    self.data['acceleration'])

        self.start = 0
        self.end = 0
        self.closestpoint = 0
        self.logger.info("Planning Ready")

        # init
        self.carx = 0
        self.cary = 0
        self.yaw = 0
        self.carvx = 0
        self.carvy = 0
        self.obss = None

    def localization_callback(self, data):
        """
        New localization Received
        """
        print("localization in")
        self.localization.CopyFrom(data)
        self.carx = self.localization.pose.position.x
        self.cary = self.localization.pose.position.y
        self.carz = self.localization.pose.position.z
        self.carvx = self.localization.pose.linear_velocity.x
        self.carvy = self.localization.pose.linear_velocity.y
        self.localization_received = True
        self.yaw = self.localization.pose.heading

    def prediction_callback(self, data):
        self.prediction.CopyFrom(data)

    def perception_callback(self, data):
        self.perception.CopyFrom(data)
        self.obss = self.perception.perception_obstacle    #list

    def routing_callback(self, data):
        self.routing.CopyFrom(data)

    def chassis_callback(self, data):
        """
        New chassis Received
        """
        self.chassis.CopyFrom(data)
        self.automode = (self.chassis.driving_mode
                         == chassis_pb2.Chassis.COMPLETE_AUTO_DRIVE)
        self.chassis_received = True

    def publish_planningmsg(self):
        """
        Generate New Path
        """
        if not self.localization_received:
            self.logger.warning(
                "localization not received yet when publish_planningmsg")
            return

        planningdata = planning_pb2.ADCTrajectory()
        now = cyber_time.Time.now().to_sec()
        planningdata.header.timestamp_sec = now
        planningdata.header.module_name = "planning"
        planningdata.header.sequence_num = self.sequence_num
        self.sequence_num = self.sequence_num + 1

        self.start = 0
        self.end = len(self.data) - 1

        planningdata.total_path_length = self.data['s'][self.end] - \
            self.data['s'][self.start]
        # self.logger.info("total number of planning data point: %d" %
        #                  (self.end - self.start))
        planningdata.total_path_time = self.data['time'][self.end] - \
            self.data['time'][self.start]
        planningdata.gear = 1
        planningdata.engage_advice.advice = \
            drive_state_pb2.EngageAdvice.READY_TO_ENGAGE

        # add dreamview path
        dns_path =  pnc_point_pb2.Path()
        dns_path.name= "dns test path"
        for i in range(self.start, self.end):
            adc_point = pnc_point_pb2.TrajectoryPoint()
            adc_point.path_point.x = self.data['x'][i]
            adc_point.path_point.y = self.data['y'][i]
            adc_point.path_point.z = self.data['z'][i]
            adc_point.v = self.data['speed'][i] * self.speedmultiplier
            adc_point.a = self.data['acceleration'][i] * self.speedmultiplier
            adc_point.path_point.kappa = self.data['curvature'][i]
            adc_point.path_point.dkappa = self.data['curvature_change_rate'][i]
            adc_point.path_point.theta = self.data['theta'][i]
            adc_point.path_point.s = self.data['s'][i]
            time_diff = self.data['time'][i] - \
                self.data['time'][0]
            adc_point.relative_time = time_diff  - (
                now - self.starttime)
            planningdata.trajectory_point.extend([adc_point])

            # path
            path_point = pnc_point_pb2.PathPoint()
            path_point.x = self.data['x'][i]+3
            path_point.y = self.data['y'][i]+1
            path_point.z = self.data['z'][i]+4

            path_point.kappa = self.data['curvature'][i]
            path_point.dkappa = self.data['curvature_change_rate'][i]
            path_point.theta = self.data['theta'][i]
            path_point.s = self.data['s'][i]
            dns_path.path_point.extend([path_point])
            planningdata.debug.planning_data.path.extend([dns_path])
        planningdata.estop.is_estop = False

        """
        dns test plot debug end
        """
        self.planning_pub.write(planningdata)
        # self.logger.debug("Generated Planning Sequence: "
        #                   + str(self.sequence_num - 1))

    def publish_planningmsg_start(self):
        """
        Generate New Path
        """
        if not self.localization_received:
            self.logger.warning(
                "localization not received yet when publish_planningmsg")
            return

        planningdata = planning_pb2.ADCTrajectory()
        now = cyber_time.Time.now().to_sec()
        planningdata.header.timestamp_sec = now
        planningdata.header.module_name = "planning"
        planningdata.header.sequence_num = self.sequence_num
        self.sequence_num = self.sequence_num + 1

        self.start = 0
        self.end = len(self.data) - 1

        planningdata.total_path_length = self.data['s'][self.end] - \
            self.data['s'][self.start]
        # self.logger.info("total number of planning data point: %d" %
        #                  (self.end - self.start))
        planningdata.total_path_time = self.data['time'][self.end] - \
            self.data['time'][self.start]
        planningdata.gear = 1
        planningdata.engage_advice.advice = \
            drive_state_pb2.EngageAdvice.READY_TO_ENGAGE

        for i in range(self.start, self.start+50):
            adc_point = pnc_point_pb2.TrajectoryPoint()
            adc_point.path_point.x = self.data['x'][i]
            adc_point.path_point.y = self.data['y'][i]
            adc_point.path_point.z = self.data['z'][i]
            adc_point.v = self.data['speed'][i] * self.speedmultiplier
            adc_point.a = self.data['acceleration'][i] * self.speedmultiplier
            adc_point.path_point.kappa = self.data['curvature'][i]
            adc_point.path_point.dkappa = self.data['curvature_change_rate'][i]
            adc_point.path_point.theta = self.data['theta'][i]
            adc_point.path_point.s = self.data['s'][i]

            time_diff = self.data['time'][i] - \
                self.data['time'][0]

            adc_point.relative_time = time_diff  - (
                now - self.starttime)

            planningdata.trajectory_point.extend([adc_point])

        planningdata.estop.is_estop = False

        self.planning_pub.write(planningdata)
        # self.logger.debug("Generated Planning Sequence: "
        #                   + str(self.sequence_num - 1))

    def publish_planningmsg_trajectory(self, trajectory):
                if not self.localization_received:
                    self.logger.warning(
                        "localization not received yet when publish_planningmsg")
                    return

                planningdata = planning_pb2.ADCTrajectory()
                now = cyber_time.Time.now().to_sec()
                planningdata.header.timestamp_sec = now
                planningdata.header.module_name = "planning"
                planningdata.header.sequence_num = self.sequence_num
                self.sequence_num = self.sequence_num + 1

                planningdata.total_path_length = self.data['s'][self.end] - \
                    self.data['s'][self.start]
                self.logger.info("total number of planning data point: %d" %
                                (self.end - self.start))
                planningdata.total_path_time = self.data['time'][self.end] - \
                    self.data['time'][self.start]
                planningdata.gear = 1
                planningdata.engage_advice.advice = \
                    drive_state_pb2.EngageAdvice.READY_TO_ENGAGE

                for i in range(len(trajectory.trajectory)-1):
                    adc_point = pnc_point_pb2.TrajectoryPoint()
                    adc_point.path_point.x = trajectory.trajectory[i][0]
                    adc_point.path_point.y = trajectory.trajectory[i][1]
                    adc_point.path_point.z = 0
                    adc_point.v =  trajectory.trajectory[i][3]
                    adc_point.a =  trajectory.trajectory[i][5]
                    adc_point.path_point.kappa = 0
                    adc_point.path_point.dkappa =0
                    adc_point.path_point.theta =  trajectory.trajectory[i][2]
                    adc_point.path_point.s = trajectory.trajectory[i][4]

                    # time_diff = self.data['time'][i] - \
                    #     self.data['time'][0]

                    time_diff = 0.1*i

                    adc_point.relative_time = time_diff  -  (now - self.starttime)
                    planningdata.trajectory_point.extend([adc_point])

                planningdata.estop.is_estop = False

                self.planning_pub.write(planningdata)
                self.logger.debug("Generated Planning Sequence: "
                                + str(self.sequence_num - 1))

    def shutdown(self):
        """
        shutdown cyber
        """
        self.terminating = True
        self.logger.info("Shutting Down...")
        time.sleep(0.2)

    def quit(self, signum, frame):
        """surrouding_obs
        shutdown the keypress thread
        """
        sys.exit(0)

    def get_obs(self):
        obs = []
        ego_obs =[self.carx, self.cary, self.carvx, self.carvy, self.yaw] #FIXME: The angle should be modified
        if self.carx == 0 and self.cary == 0:
            return []
        # print("ego_yaw",self.carx, self.cary, self.yaw)
        obs.append(ego_obs)
        if self.obss is not None:
            for i in range(len(self.obss)):
                obs_info = self.obss[i]
                surrouding_obs = [obs_info.position.x,
                                                    obs_info.position.y,
                                                    obs_info.velocity.x,
                                                    obs_info.velocity.y,
                                                    obs_info.theta]
                obs.append(surrouding_obs)
        return obs

class Werling_planner_SP():
    def  __init__(self,):
        self.trajectory_planner = JunctionTrajectoryPlanner_SP_Dns()
        self.dynamic_map = DynamicMap()
        self.read_ref_path_from_file()

    def update_path(self, obs, done):# TODO: represent a obs
        if done or len(obs)<1:
            print("Without Obs")
            self.trajectory_planner.clear_buff(clean_csp=False)
        else:

            self.dynamic_map.update_map_from_list_obs(obs)

            trajectory_action, index = self.trajectory_planner.trajectory_update(self.dynamic_map)
            chosen_action_id = index

            chosen_trajectory = self.trajectory_planner.trajectory_update_CP(chosen_action_id)

            # print("chosen_",chosen_trajectory.trajectory)
            return chosen_trajectory

    def read_ref_path_from_file(self):
        record_file = os.path.join(APOLLO_ROOT, 'data/log/garage.csv')
        try:
            file_handler = open(record_file, 'r')
        except (FileNotFoundError, IOError) as ex:
            self.logger.error("Error opening {}: {}".format(record_file, ex))
            sys.exit(1)

        self.data = genfromtxt(file_handler, delimiter=',', names=True)
        file_handler.close()
        t_array = []
        self.ref_path = Lane()

        for i in range(0,len(self.data)//100): # The Apollo record data is too dense!
            lanepoint = Lanepoint()
            lanepoint.position.x = self.data['x'][i*90]
            lanepoint.position.y = self.data['y'][i*90]
            # print("ref path", lanepoint.position.x, lanepoint.position.y)
            self.ref_path.central_path.append(lanepoint)
            t_array.append(lanepoint)
        self.ref_path.central_path_array = np.array(t_array)
        self.ref_path.speed_limit = 60/3.6 # m/s
        self.dynamic_map.update_ref_path_from_routing(self.ref_path)

def main():
    """
    Main cyber
    """

    node = cyber.Node("rtk_player")

    Logger.config(
        log_file=os.path.join(APOLLO_ROOT, 'data/log/rtk_player.log'),
        use_stdout=True,
        log_level=logging.DEBUG)

    record_file = os.path.join(APOLLO_ROOT, 'data/log/garage.csv')

    player = RtkPlayer(record_file, node)

    planner = Werling_planner_SP()

    atexit.register(player.shutdown)

    node.create_reader('/apollo/canbus/chassis', chassis_pb2.Chassis,
                       player.chassis_callback)

    node.create_reader('/apollo/localization/pose',
                       localization_pb2.LocalizationEstimate,
                       player.localization_callback)

    node.create_reader('/apollo/prediction',
                        prediction_obstacle_pb2.PredictionObstacles,
                        player.prediction_callback)

    node.create_reader('/apollo/perception/obstacles',
                        perception_obstacle_pb2.PerceptionObstacles,
                        player.perception_callback)

    RTK = 0

    while not cyber.is_shutdown():
        now = cyber_time.Time.now().to_sec()
        # # New add
        if not RTK:
            obs =  player.get_obs()
            trajectory = planner.update_path(obs, done=0)
            if trajectory is not None:
                player.publish_planningmsg_trajectory(trajectory)
        else:
            # player.publish_planningmsg()
            player.publish_planningmsg_start()
        sleep_time = 0.1 - (cyber_time.Time.now().to_sec() - now)
        if sleep_time > 0:
            time.sleep(sleep_time)

if __name__ == '__main__':
    cyber.init()
    main()

    cyber.shutdown()


import time
import jwt
from backlogic.models import NurseInfo,PatientInfo

secret1 = '1919810'
secret2 = '紫者，妪也，紫大寿，强宴宾客于迷途家。肴即尽，取陈酒，觞诸宾。宾皆不敢，尽却之。唯天子许，饮三壶，大醉。紫甚欣慰，请起舞悦众宾。众思谢之恐死，苦不敢言。遂舞。笙箫起，笛音生，紫披发旖旎，自思艳绝。舞毕，遽顾问于客。客皆骇然称善。唯天子醉曰：“笑煞人，老妪何惺惺然作处子态!”鞭于庭。'
def generate_jwt(payload, expiry=None):
    """

    :param payload: dict 载荷
    :param expiry: datetime 有效期
    :return: 生成jwt
    """
    if expiry is None:
        expiry = int(time.time()) + 60 *5

    _payload = {'exp': expiry}
    _payload.update(payload)
    #token = ''
    token = jwt.encode(_payload, secret1, algorithm='HS256')
    return token

def decode_jwt(token):
    payload = jwt.decode(token,secret2,algorithms=['HS256'])
    #payload = ''
    return payload

def decode_jwt_back(token):
    try:
        payload = jwt.decode(token,secret1,algorithms=['HS256'])
        return payload
    except jwt.exceptions.ExpiredSignatureError:
        return False

def verify_jwt(token):
    obj = NurseInfo.objects.filter(login_jwt = token)
    try:
        if obj[0].workPermitNumber:
            # print( obj[0].login_jwt)
            return obj[0].workPermitNumber
        else:
            print('工作证号不存在')
            return False
    except IndexError:
        return False

def remoteLoginVerification(request):
    # 获取该名护士照顾的所有病人的信息
    # 验证http请求head中的jwt
    verification_jwt = request.META['HTTP_LOGINJWT']  #规定jwt存在header的HTTP_LOGINJWT中
    print('patient2nurse:')
    print(verification_jwt)

    # 建议传过来的jwt是否可解码
    payload = decode_jwt_back(verification_jwt)
    if payload:
        pass
    else:
        return false
        # returnjson = {'state':'400','message':'登录状态异常,jwt不存在'}
        # return HttpResponse(json.dumps(returnjson))

    # 可解码的情况下，判断该jwt是否与数据库中的jwt相同，不同则说明这个jwt是被顶掉的
    idnurse = verify_jwt(verification_jwt)
    print(idnurse)
    if idnurse:
        return true
    else:
        return false


#!/usr/bin/env python

"""
"""

from maya import cmds
from crefor import api
from crefor.lib import libName
from crefor import log

import unittest
import logging

logger = log.get_logger(__name__)
logger.setLevel(logging.CRITICAL)
logger.disabled = True

class TestApi(unittest.TestCase):

    def __create(self):
        """
        """

        arm = api.create("L", "arm", 0)
        spine = api.create("C", "spine", 0)

        return arm, spine

    def setUp(self):
        """Runs before each test"""
        cmds.file(newFile=True, force=True)

    def tearDown(self):
        """Runs after each test"""
        pass

    def test_create(self):
        """
        Test api.create()
        """

        arm, _ = self.__create()
        self.assertEquals(api.exists(arm),
                          True,
                          "Guide does not exist: %s" % arm.node)

    def test_set_parent(self):
        """
        Test api.set_parent(child, parent)
        """

        child, parent = self.__create()

        api.set_parent(child, parent)

        self.assertEquals(child.parent.node,
                          parent.node,
                          "'%s' is not parent of '%s'" % (parent.node, child.node))
        self.assertEquals(child.node in [g.node for g in parent.children],
                          True,
                          "'%s' is not a child of '%s'" % (child.node, parent.node))

    def test_add_child(self):
        """
        Test api.add_child(parent, child)
        """

        child, parent = self.__create()

        api.add_child(parent, child)

        self.assertEquals(child.parent.node,
                          parent.node,
                          "'%s' is not parent of '%s'" % (parent.node, child.node))
        self.assertEquals(child.node in [g.node for g in parent.children],
                          True,
                          "'%s' is not a child of '%s'" % (child.node, parent.node))

    def test_has_parent(self):
        """
        Test api.has_parent(parent, child)
        """

        child, parent = self.__create()

        api.set_parent(child, parent)

        self.assertEquals(api.has_parent(child, parent),
                          True,
                          "'%s' is not parent of '%s'" % (parent.node, child.node))

    def test_has_child(self):
        """
        Test api.has_child(parent, child)
        """

        child, parent = self.__create()

        api.set_parent(child, parent)

        self.assertEquals(api.has_child(parent, child),
                          True,
                          "'%s' is not parent of '%s'" % (parent.node, child.node))

    def test_is_parent(self):
        """
        Test api.has_child(parent, child)
        """

        child, parent = self.__create()

        api.set_parent(child, parent)

        self.assertEquals(api.is_parent(parent, child),
                          True,
                          "'%s' is not parent of '%s'" % (parent.node, child.node))

    def test_remove(self):
        """
        Test api.remove(guide)
        """

        child, _ = self.__create()

        self.assertEquals(child.exists(),
                          True,
                          "Guide '%s' does not exist." % child.node)
        child.remove()

        self.assertEquals(child.exists(),
                          False,
                          "Guide '%s' does not exist." % child.node)

    def test_remove_parent(self):
        """
        Test api.remove_parent(guide)
        """

        child, parent = self.__create()

        api.set_parent(child, parent)

        api.remove_parent(child)

        self.assertEquals(api.is_parent(parent, child),
                          False,
                          "'%s' is not parent of '%s'" % (parent.node, child.node))

    def test_duplicate(self):
        """
        Test api.duplicate(guide)
        """

        child, _ = self.__create()

        self.assertEquals(api.exists(child),
                          True,
                          "Guide does not exist: %s" % child.node)

        duplicate = api.duplicate(child)

        self.assertEquals(child < duplicate and duplicate > child,
                          True,
                          "'%s' is not a duplicate of '%s'" % (duplicate, child))

    def test_compile(self):
        """
        Test api.compile()
        """

        child, _ = self.__create()

        self.assertEquals(api.exists(child),
                          True,
                          "Guide does not exist: '%s'" % child.node)

        api.compile()

        self.assertEquals(api.exists(child),
                          False,
                          "Guide exists: '%s'" % child.node)

        joint = libName.update(child.node, suffix="jnt")

        self.assertEquals(cmds.objExists(joint),
                          True,
                          "Joint does not exist: '%s'" % joint)

        guides = api.get_guides()
        self.assertEquals(guides == [],
                          True,
                          "Api listed guides after compiling: %s" % guides)

    def test_get_guides(self):
        """
        Test api.get_guides()
        """

        child, parent = self.__create()

        guides = api.get_guides()

        self.assertEquals(len(guides) == 2,
                          True,
                          "Failed to find all guides: '%s'" % guides)

        self.assertEquals(child in guides and parent in guides,
                          True,
                          "Created nodes are not listed: %s" % guides)

class TestApiReinit(TestApi):

    def __create(self):
        """
        """

        arm = api.create("L", "arm", 0)
        spine = api.create("C", "spine", 0)

        arm = arm.reinit()
        spine = spine.reinit()

        return arm, spine


from typing import Optional
from fastapi import Query, HTTPException
from pymatgen.core.periodic_table import Element
from pymatgen.core import Composition
from maggma.api.query_operator import QueryOperator
from maggma.api.utils import STORE_PARAMS

from collections import defaultdict

class MoleculeElementsQuery(QueryOperator):
    """
    Method to generate a query on molecules using a list of elements
    """

    def query(
        self,
        elements: Optional[str] = Query(
            None,
            description="Query by elements in the material composition as a comma-separated list",
        ),
    ) -> STORE_PARAMS:
        crit = {}

        if elements:
            try:
                element_list = [Element(e.strip()) for e in elements.split(",")]
            except ValueError:
                raise HTTPException(
                    status_code=400,
                    detail="Problem processing one or more provided elements.",
                )
            crit["elements"] = {"$all": [str(el) for el in element_list]}

        return {"criteria": crit}

    def ensure_indexes(self):  # pragma: no cover
        return [("elements", False)]

class MoleculeBaseQuery(QueryOperator):
    """
    Method to generate a query over molecule data.
    """

    def query(
        self,
        nelements_max: Optional[float] = Query(
            None,
            description="Maximum value for the number of elements.",
        ),
        nelements_min: Optional[float] = Query(
            None,
            description="Minimum value for the number of elements.",
        ),
        EA_max: Optional[float] = Query(
            None,
            description="Maximum value for the electron affinity in eV.",
        ),
        EA_min: Optional[float] = Query(
            None,
            description="Minimum value for the electron affinity in eV.",
        ),
        IE_max: Optional[float] = Query(
            None,
            description="Maximum value for the ionization energy in eV.",
        ),
        IE_min: Optional[float] = Query(
            None,
            description="Minimum value for the ionization energy in eV.",
        ),
        charge_max: Optional[int] = Query(
            None,
            description="Maximum value for the charge in +e.",
        ),
        charge_min: Optional[int] = Query(
            None,
            description="Minimum value for the charge in +e.",
        ),
        pointgroup: Optional[str] = Query(
            None,
            description="Point of the molecule in Schoenflies notation.",
        ),
        smiles: Optional[str] = Query(
            None,
            description="The simplified molecular input line-entry system (SMILES) \
            representation of the molecule.",
        ),
    ) -> STORE_PARAMS:
        crit = defaultdict(dict)  # type: dict

        d = {
            "nelements": [nelements_min, nelements_max],
            "EA": [EA_min, EA_max],
            "IE": [IE_min, IE_max],
            "charge": [charge_min, charge_max],
        }  # type: dict

        for entry in d:
            if d[entry][0] is not None:
                crit[entry]["$gte"] = d[entry][0]

            if d[entry][1] is not None:
                crit[entry]["$lte"] = d[entry][1]

        if pointgroup:
            crit["pointgroup"] = pointgroup

        if smiles:
            crit["smiles"] = smiles

        return {"criteria": crit}

    def ensure_indexes(self):  # pragma: no cover
        keys = self._keys_from_query()
        indexes = []
        for key in keys:
            if "_min" in key:
                key = key.replace("_min", "")
                indexes.append((key, False))
        return indexes

class MoleculeFormulaQuery(QueryOperator):
    """
    Method to generate a query for molecule data using a chemical formula
    """

    def query(
        self,
        formula: Optional[str] = Query(
            None,
            description="Chemical formula of the molecule. A comma-separated list of formulas is also accepted.",
        ),
    ) -> STORE_PARAMS:
        crit = defaultdict(dict)  # type: dict

        if formula:
            formula_list = [f.strip() for f in formula.split(",")]

            try:
                if len(formula_list) == 1:
                    reduced_formula = Composition(
                        formula
                    ).get_reduced_formula_and_factor()[0]
                    crit["formula_pretty"] = reduced_formula
                else:
                    crit["formula_pretty"] = {
                        "$in": [
                            Composition(f).get_reduced_formula_and_factor()[0]
                            for f in formula_list
                        ]
                    }
            except ValueError:
                raise HTTPException(
                    status_code=400,
                    detail="Problem processing one or more provided formulas.",
                )

        return {"criteria": crit}

    def ensure_indexes(self):  # pragma: no cover
        return [("formula_pretty", False)]


from .asn1 import Decoder, Encoder, Number

class Package:
    pdu_id = None
    version = None
    community = None
    pdu = None

    def __init__(self):
        self.request_id = None
        self.error_status = None
        self.error_index = None
        self.variable_bindings = []

    def encode(self):
        encoder = Encoder()

        with encoder.enter(Number.Sequence):
            encoder.write(self.version, Number.Integer)
            encoder.write(self.community, Number.OctetString)

            self.pdu.request_id = self.request_id
            self.pdu.encode(encoder)

        return encoder.output()

    def decode(self, data):
        decoder = Decoder(data)
        with decoder.enter():
            decoder.read()  # version
            decoder.read()  # community

            with decoder.enter():
                _, self.request_id = decoder.read()
                _, self.error_status = decoder.read()
                _, self.error_index = decoder.read()

                with decoder.enter():
                    while not decoder.eof():
                        with decoder.enter():
                            _, oid = decoder.read()
                            tag, value = decoder.read()
                            self.variable_bindings.append((oid, tag, value))

class SnmpMessage(Package):

    @classmethod
    def make(cls, version, community, pdu):
        pkg = cls()
        pkg.version = version
        pkg.community = community
        pkg.pdu = pdu
        return pkg


from typing import Any
from PIL import Image
from torchvision import transforms
from torchvision.datasets import VisionDataset
import utils.ss_transforms as tr
import matplotlib.pyplot as plt
from torch import from_numpy
import numpy as np
from matplotlib.colors import ListedColormap

class_map={
    1: 13,  # ego_vehicle : vehicle
   7: 0,   # road
   8: 1,   # sidewalk
   11: 2,  # building
   12: 3,  # wall
   13: 4,  # fence
   17: 5,  # pole
   18: 5,  # poleGroup: pole
   19: 6,  # traffic light
   20: 7,  # traffic sign
   21: 8,  # vegetation
   22: 9,  # terrain
   23: 10,  # sky
   24: 11,  # person
   25: 12,  # rider
   26: 13,  # car : vehicle
   27: 13,  # truck : vehicle
   28: 13,  # bus : vehicle
   32: 14,  # motorcycle
   33: 15,  # bicycle
}
palette=np.array([
    [128,64,128],       #road
    [244,35,232],       #sidewalk
    [70,70,70],         #building
    [102,102,156],      #wall
    [190,153,153],      #fence
    [153,153,153],      #pole
    [250,170,30],       #light
    [220,220,0],        #sign
    [107,142,35],       #vegetation
    [152,251,152],      #terrain
    [70,130,180],       #sky
    [220,20,60],        #person
    [255,0,0],          #rider
    [0,0,142],          #vehicle
    [0,0,230],          #motorcycle
    [119,11,32],        #bycicle
    [0,0,0]] )
class GTAVDataset(VisionDataset):

    def __init__(self,
                 root: str,
                 list_samples: list[str]=[],
                 fileName:str=None,
                 transform: tr.Compose = None,
                 client_name: str = None):
        super().__init__(root=root, transform=transform, target_transform=None)

        samples_from_file=[]
        if fileName!=None:
            f = open(self.root+'/'+fileName, "r")
            samples_from_file=[x for x in f.read().split('\n') if x != '']
        self.list_samples = list_samples+samples_from_file
        self.client_name = client_name
        self.target_transform = self.get_mapping()

    @staticmethod
    def get_mapping():
        map_classes=np.zeros((256,),dtype=np.int64)
        for i in range(256):
            map_classes[i]=class_map[i] if i in class_map else 255
        return lambda x: from_numpy(map_classes[x])

    def __getitem__(self, index: int) -> Any:
        image=Image.open(self.root+'/images/'+self.list_samples[index])
        target=Image.open(self.root+'/labels/'+self.list_samples[index])

        if self.transform:
            image,target=self.transform(image,target)
        if self.target_transform:
            target=self.target_transform(target)
        return image, target

    def __len__(self) -> int:
        return len(self.list_samples)

    def showSample(self,index=0,prediction=None):
        cMap=ListedColormap(palette/256)
        image,label=self[index]
        k=3
        if prediction==None:
            k=2
        fig=plt.figure(figsize=(10*k*1080/1920,10))
        fig.add_subplot(1,k,1)
        plt.imshow(image.permute(1,2,0))
        plt.axis('off')
        fig.add_subplot(1,k,2)
        plt.imshow(label,vmax=16,cmap=cMap,interpolation='none')
        plt.axis('off')
        if prediction==None:
            return
        fig.add_subplot(1,3,3)
        plt.imshow(prediction,vmax=16,cmap=cMap,interpolation='none')
        plt.axis('off')