import numpy as np
from osgeo import ogr, osr
# import json, sys

import time
from topoflow.utils import shape_utils as su

#---------------------------------------------------------------------
def get_basin_repo_dir():

    #-----------------------------------
    # Modify this directory as needed.
    #-----------------------------------
    repo_dir  = '/Users/peckhams/Dropbox/NOAA_NextGen/'
    repo_dir += '__NextGen_Example_Basin_Repo/'
    return repo_dir

#   get_basin_repo_dir()
#---------------------------------------------------------------------
def get_camels_data_dir(version='1.2'):

    if (version == '1.2'):
       camels_dir = '/Users/peckhams/Data/CAMELS/2016_v1.2/'
    elif (version == '2.0'):
       camels_dir = '/Users/peckhams/Data/CAMELS/2023_v2.0/'
    else:
       repo_dir = get_basin_repo_dir()
       camels_dir = repo_dir + 'CAMELS/'

    return camels_dir

#   get_camels_data_dir()
#---------------------------------------------------------------------
def create_camels_bbox_file( bbox_file='camels_bbox.txt',
                             nb_max=700):

    #-------------------------------------------------------------
    # This function creates a semicolon-delimited text file,
    # similar to other camels files (e.g. camels_topo.txt),
    # that contains gauge_id, minlon, maxlon, minlat, & maxlat.
    # It reads from a merged shapefile that has the basin
    # boundaries for all HRUs within any given CAMELS basin.
    #-------------------------------------------------------------
    # Change these directories as needed.
    # Set nb_max to smaller value (e.g. 20) for testing.
    #-------------------------------------------------------------
    camels1_dir = get_camels_data_dir(version='1.2')
    camels2_dir = get_camels_data_dir(version='2.0')
    #-------------------------------------------------------------
    shape_dir   = camels1_dir + 'shapefiles/merge/'
    shape_path  = shape_dir + 'basinset_gf_nhru.shp'
    prj_path    = shape_dir + 'basinset_gf_nhru.prj'
    #-------------------------------------------------------------
    bbox_unit = open( camels2_dir + bbox_file, 'w' )

    start_time = time.time()
    shape_unit = ogr.Open( shape_path )
    layer      = shape_unit.GetLayer()
    n_hrus     = np.int32(0)
    n_basins   = np.int32(0)
    n_hrus_in_basin = np.int32(0)
    SWAP_XY    = True     #######
    delim      = ';'  # semi-colon delimiter
    last_gauge_id = '00000000'
    minlon = 180.0
    maxlon = -180.0
    minlat = 90.0
    maxlat = -90.0

    #---------------------------------
    # Write header for new bbox_file
    #---------------------------------
    ## new_headings=['gauge_id','MINLON','MAXLON','MINLAT','MAXLAT']
    bbox_unit.write('gauge_id' + delim)
    bbox_unit.write('MINLON'   + delim)
    bbox_unit.write('MAXLON'   + delim)
    bbox_unit.write('MINLAT'   + delim)
    bbox_unit.write('MAXLAT'   + '\n')

    #-----------------------------------------
    # Iterate over the features in the layer
    # GeometryType = 3 is POLYGON.
    #-----------------------------------------
    print('Working...')
    for feature in layer:
        geometry   = feature.GetGeometryRef()
        attributes = feature.items()  # This is a Python dictionary
        # n_rings    = geometry.GetGeometryCount()

        gauge_id = attributes['GAGEID']
        FIRST_ONE = (last_gauge_id == '00000000')
        # FIRST_ONE = (n_hrus == 0)

        #----------------------------------
        # Get all points in this geometry
        #----------------------------------
        x1, y1 = su.get_polygon_points( feature )
#         print('x1 =', x1)
#         print('y1 =', y1)
        # break

        #------------------------------------------
        # Convert coords to Geo. lon/lat (WGS-84)
        #------------------------------------------
        x2,y2 = su.convert_coords(x1,y1, inPRJfile=prj_path,
                                  SWAP_XY=SWAP_XY, PRINT=False)
#         print('x2 =', x2)
#         print('y2 =', y2)
#         break

        #--------------------------------------
        # Get geographic bounding box for one
        # of the HRUs in this CAMELS basin
        #--------------------------------------
        hru_minlon, hru_maxlon, hru_minlat, hru_maxlat = su.get_bounding_box(x2, y2)
#         print()
#         print('minlon, maxlon =', minlon, maxlon)
#         print('minlat, maxlat =', minlat, maxlat)
        n_hrus += 1

        if ((gauge_id == last_gauge_id) or FIRST_ONE):
            minlon = min(hru_minlon, minlon)
            maxlon = max(hru_maxlon, maxlon)
            minlat = min(hru_minlat, minlat)
            maxlat = max(hru_maxlat, maxlat)
        else:
            #-------------------------------------
            # Write line of info to new CSV file
            #---------------------------------------------------------
            # su.get_bounding_box() rounds to 5 digits after decimal
            #---------------------------------------------------------
            bbox_unit.write( gauge_id    + delim )
            bbox_unit.write( str(minlon) + delim )
            bbox_unit.write( str(maxlon) + delim )
            bbox_unit.write( str(minlat) + delim )
            bbox_unit.write( str(maxlat) + '\n' )
            n_basins += 1
            print('n_basins =', n_basins, 'of 671.')
            print('   gauge_id =', gauge_id)
            print('   n_hrus_in_basin =', n_hrus_in_basin)
            n_hrus_in_basin = np.int32(0)
            if (n_basins >= nb_max):
                break
            #----------------------------
            # Reset bounding box coords
            #----------------------------
            minlon = hru_minlon
            maxlon = hru_maxlon
            minlat = hru_minlat
            maxlat = hru_maxlat

        n_hrus_in_basin += 1
        last_gauge_id = gauge_id

    #--------------------------------
    # Print info, close files, etc.
    #--------------------------------
    print('Processed', n_hrus, 'HRUs.')
    print('Processed', n_basins, 'CAMELS basins.')
    bbox_unit.close()
    run_time = (time.time() - start_time)
    print('Run time =', run_time, '[sec]')
    print('Finished.')

#   create_camels_bbox_file()
#---------------------------------------------------------------------
def sort_camels_bbox_file( bbox_file='camels_bbox0.txt',
         bbox_file2='camels_bbox.txt'):

    print('Reading lines from:', bbox_file)
    camels2_dir = get_camels_data_dir(version='2.0')
    bbox_unit = open( camels2_dir + bbox_file, 'r' )
    lines    = bbox_unit.readlines()
    headings = lines[0]
    lines    = lines[1:]
    bbox_unit.close()
    print('len(lines) =', len(lines))

    #--------------------------------------
    # Write sorted lines to new bbox file
    #--------------------------------------
    print('Writing sorted lines to:', bbox_file2)
    lines.sort()
    bbox_unit2 = open( camels2_dir + bbox_file2, 'w' )
    bbox_unit2.write(headings)
    bbox_unit2.writelines( lines )
    bbox_unit2.close()
    print('Finished.')
    print()

#   sort_camels_bbox_file()
#---------------------------------------------------------------------
def find_extra_ids( CHECK_DUPES=False ):

    #--------------------------------------------------------------
    # Note: camels_bbox.txt has 678 lines, while camels_name.txt
    #       and the others only have 672 lines.  This function
    #       helps to find the extra gauge IDs.
    #--------------------------------------------------------------
    # Note: In the shapefile attribute table, the gauge ID:
    #       07067000 does occur first in a set of 77 HRUs that
    #       have hru_id_loc in {1,...,77}, and then in another
    #       set of 10 HRUs that have hru_id_loc in {1,...,10}.
    #       Some have gtype="polygon" and some have gtype=
    #       "multipolygon".
    #--------------------------------------------------------------
    # Results of this function:
    #
    # gauge IDs in bbox file not in name file:
    # ['06775500', '01150900', '06846500',
    # '09535100', '02081113', '03448942']
    #
    # gauge IDs in name file not in bbox file:
    # ['01333000']
    #
    # 01150900;-73.00165;-72.33164;43.64477;44.12633
    # 0208111310;-89.60984;-89.24013;30.44601;31.10538
    # 0344894205;-87.88285;-87.21651;35.18618;35.58351
    # 06775500;-105.93089;-105.83394;40.49394;40.56312
    # 06846500;-96.48799;-96.06662;39.86638;40.34738
    # 09535100;-112.1861;-111.8549;33.91465;34.11396
    #
    # 01333000;02;GREEN RIVER AT WILLIAMSTOWN, MA
    #--------------------------------------------------------------
    # Duplicate ID in bbox file (which one is right?):
    # One has 77 hrus, but they don't get highlighted in
    # QGIS with View > Zoom to Selection.  Other has 10 hrus.
    #
    # 07067000;-92.02911;-91.3148;36.94448;37.25004 (wider)
    # 07067000;-94.73446;-94.31482;32.74102;33.23056 (taller)
    #
    # 07067000;11;Current River at Van Buren, MO
    # See USGS website for this river:
    # https://waterdata.usgs.gov/monitoring-location/
    #    07067000/#parameterCode=00065&period=P7D
    #--------------------------------------------------------------
    # In QGIS, install the BoundingBox plugin (how to use?).
    # And/or right click on a point & select Copy Coordinate.
    #--------------------------------------------------------------
    camels2_dir = get_camels_data_dir(version='2.0')
    # bbox_file = camels2_dir + 'camels_bbox_677.txt'
    bbox_file = camels2_dir + 'camels_bbox_672.txt'
    name_file = camels2_dir + 'camels_name.txt'

    bbox_unit = open( bbox_file, 'r' )
    lines1 = bbox_unit.readlines()
    bbox_unit.close()
    print('Number of IDs in bbox_file =', len(lines1)-1)
    print('Number of unique IDs in bbox_file =', len(set(lines1))-1)
    print()

    name_unit = open( name_file, 'r' )
    lines2 = name_unit.readlines()
    name_unit.close()
    print('Number of IDs in name_file =', len(lines2)-1)
    print('Number of unique IDs in name_file =', len(set(lines2))-1)

    k = 0
    id_list1 = list()
    for line in lines1:
        id = lines1[k][:8]
        id_list1.append( id )
        k += 1
        # print('id =', id)

    k = 0
    id_list2 = list()
    for line in lines2:
        id = lines2[k][:8]
        id_list2.append( id  )
        k += 1

    if (CHECK_DUPES):
        # This requires that id_list1 is sorted.
        print('Checking for duplicates...')
        last_id = '00000000'
        for id in id_list1:
            if (id == last_id):
                print('duplicate id =', id)
            last_id = id

#     k = 0
#     for id1 in id_list1[:-1]:
#         id2 = id_list2[k]
#         if (id2 != id1):
#             print('id1, id2 =', id1, id2)
#         k += 1

    #---------------------------
    # Compare the two ID lists
    #---------------------------
    print()
    diff1 = list( set(id_list1) - set(id_list2))
    print('gauge IDs in bbox file not in name file:')
    print(diff1)
    print()

    diff2 = list( set(id_list2) - set(id_list1))
    print('gauge IDs in name file not in bbox file:')
    print(diff2)
    print()

    print('Finished.')

#   find_extra_ids()
#---------------------------------------------------------------------
def create_tsv( new_tsv_file = 'new_camels.tsv', delim='\t' ):

    #--------------------------------------------------------------
    # Note: CAMELS version 1.2 was released in 2016 and has much
    #       more data including basin shapefiles. (91 GB)
    #       CAMELS version 2.0 was released in July 2023 and
    #       does not include basin shapefiles. (7.5 GB)
    #--------------------------------------------------------------
    #       Change these directories as needed.
    #--------------------------------------------------------------
    camels1_dir = get_camels_data_dir(version='1.2')
    camels2_dir = get_camels_data_dir(version='2.0')
    tsv_dir     = get_camels_data_dir(version='repo')

    #---------------------------------------------------
    # These CAMELS data files use ";" as the delimeter
    # All files start with "gauge_id" and are sorted
    #   by gauge_id.  All gauge IDs are 8 chars.
    #---------------------------------------------------
    name_file   = 'camels_name.txt'  # all in camels2_dir
    bbox_file   = 'camels_bbox_671.txt'
    topo_file   = 'camels_topo.txt'
    clim_file   = 'camels_clim.txt'
    soil_file   = 'camels_soil.txt'
    geol_file   = 'camels_geol.txt'

    #--------------------------------------------------
    # Open text files to read & new csv file to write
    #--------------------------------------------------
    name_unit = open( camels2_dir + name_file, 'r' )
    bbox_unit = open( camels2_dir + bbox_file, 'r' )
    topo_unit = open( camels2_dir + topo_file, 'r' )
    clim_unit = open( camels2_dir + clim_file, 'r' )
    # soil_unit = open( camels2_dir + soil_file, 'r' )
    # geol_unit = open( camels2_dir + geol_file, 'r' )
    #---------------------------------------------------
    tsv_unit = open( tsv_dir + new_tsv_file, 'w' )

    #---------------------------
    # Read the column headings
    #---------------------------
    name_line = name_unit.readline()
    name_head = name_line.replace(';', delim)
    name_head = name_head.replace('\n', '')
    #-------------------------------------------------
    bbox_line = bbox_unit.readline()
    bbox_head = bbox_line.replace(';', delim)
    bbox_head = bbox_head[9:]   # exclude "gauge_id"
    bbox_head = bbox_head.replace('\n', '')
    #-------------------------------------------------
    topo_line = topo_unit.readline()
    topo_head = topo_line.replace(';', delim)
    topo_head = topo_head[9:]   # exclude "gauge_id"
    topo_head = topo_head.replace('\n', '')
    #-------------------------------------------------
    clim_line = clim_unit.readline()
    clim_head = clim_line.replace(';', delim)
    clim_head = clim_head[9:]   # exclude "gauge_id"
    clim_head = clim_head.replace('\n', '')

    #----------------------------------------
    # Write column headings to new tsv file
    #----------------------------------------
    tsv_unit.write( name_head + delim )
    tsv_unit.write( bbox_head + delim )
    tsv_unit.write( topo_head + delim )
    tsv_unit.write( clim_head + '\n')

    #-----------------------------------------
    # Don't forget about minlon,maxlon, etc.      ##############
    # Add afterwards w/ separate function?
    #-----------------------------------------

    #-----------------------------
    # Write data to new tsv file
    #--------------------------------------------
    # Note: All gauge IDs are 8 chars.
    # Note: "gauge_id" heading is also 8 chars.
    # Note: All files are sorted by gauge ID.
    # Note: "readline" result ends w/ newline.
    #--------------------------------------------
    print('Working...')
    while (True):
        name_line = name_unit.readline()
        if (name_line == ''):
            break  # (reached end of file)
        name_line = name_line.replace(';', delim)
        tsv_unit.write( name_line[:-1] + delim )
        #-------------------------------------------
        bbox_line = bbox_unit.readline()
        bbox_line = bbox_line.replace(';', delim)
        tsv_unit.write( bbox_line[9:-1] + delim )
        #-------------------------------------------
        topo_line = topo_unit.readline()
        topo_line = topo_line.replace(';', delim)
        tsv_unit.write( topo_line[9:-1] + delim )
        #-------------------------------------------
        clim_line = clim_unit.readline()
        clim_line = clim_line.replace(';', delim)
        tsv_unit.write( clim_line[9:-1] + '\n')    # notice newline
        #-------------------------------------------
#         soil_line = soil_unit.readline()
#         soil_line = soil_line.replace(';', delim)
#         tsv_unit.write( soil_line[9:-1] + delim)
#         #-------------------------------------------
#         geol_line = geol_unit.readline()
#         geol_line = geol_line.replace(';', delim)
#         tsv_unit.write( geol_line[9:-1] + '\n')

    #------------------
    # Close all files
    #------------------
    print('Finished.')
    name_unit.close()
    bbox_unit.close()
    topo_unit.close()
    clim_unit.close()
    # soil_unit.close()
    # geol_unit.close()
    #--------------------
    tsv_unit.close()

#   create_tsv()
#---------------------------------------------------------------------


from django.conf.urls import patterns, include, url
from taskapp import  views
from django.contrib import admin
admin.autodiscover()

urlpatterns = patterns('',
    # Examples:
    #url(r'^$', 'task_project.views.home', name='home'),
    # url(r'^blog/', include('blog.urls')),

    url(r'^admin/', include(admin.site.urls)),
    url(r'^tasks/?$', views.index),
    url(r'^addtask', views.new_task, name="newTask"),
    url(r'^register', views.register, name='register'),
    #url(r'^newUser', views.new_user, name='newUSer'),
    url(r'^login', views.user_login, name='login'),
    url(r'^$', views.user_login, name='login'),

)


import argparse
import os
import re
import json
import gzip
from collections import defaultdict

os.environ['HF_DATASETS_OFFLINE'] = '1'  # ask datasets library not to make arbitrary web requests
os.environ['TRANSFORMERS_OFFLINE'] = '1'  # ask transformer library not to make arbitrary web requests
os.environ['BITSANDBYTES_NOWELCOME'] = '1'  # disable welcome message that bitsandbytes prints, it's unnecessary noise

import numpy as np
import torch
import peft
import transformers
import trl

DATA_JSON_FILES_PATH = "./huggingface_cache/datasets--Anthropic--hh-rlhf/snapshots/09be8c5bbc57cb3887f3a9732ad6aa7ec602a1fa/helpful-online/"  # downloaded by running `make download-datasets-and-models` in this repo
BASE_MODEL_PATH = "./llama2-supervised-finetuning-output/final_checkpoint_merged"  # generated as the output of running `make train-supervised-finetuning`
OUTPUT_DIRECTORY = "./llama2-direct-preference-finetuning-output"
RANDOMNESS_SEED = 0
BATCH_SIZE = 1  # number of samples seen per gradient update - to increase training speed, set this to the largest size that your hardware can support without running out of memory
CONTEXT_WINDOW_SIZE = 1300  # maximum length of any input to the model, used to filter out too-long data points (this doesn't have to be the same value as in supervised_finetuning.py) - to improve performance on longer prompts, set this to the largest size that your hardware can support without running out of memory
TRAINING_STEPS = 1000  # number of steps to train for (since we're using gradient_accumulation_steps=4, the model will see TRAINING_STEPS * 4 * BATCH_SIZE samples throughout the entire training run)

def prompt_formatter(example):
    chosen_messages, rejected_messages = re.split(r"\n\n(Human|Assistant): ", example["chosen"])[1:], re.split(r"\n\n(Human|Assistant): ", example["rejected"])[1:]
    chosen_prompt_so_far = ""
    if len(chosen_messages) != len(rejected_messages) or len(chosen_messages) % 4 != 0:
        return []
    result = []
    for i in range(0, min(len(chosen_messages), len(rejected_messages)), 4):
        if not (chosen_messages[i] == rejected_messages[i] == "Human"):
            return []
        if not (chosen_messages[i + 2] == rejected_messages[i + 2] == "Assistant"):
            return []
        prompt, chosen, rejected = chosen_messages[i + 1], chosen_messages[i + 3], rejected_messages[i + 3]
        chosen_prompt_so_far += f'\n\n### Human:\n{prompt}\n\n### Assistant:\n'
        result.append({"prompt": chosen_prompt_so_far, "chosen": chosen, "rejected": rejected})
        chosen_prompt_so_far += chosen_messages[i + 3]
    return result

def create_datasets(tokenizer):
    with gzip.open(os.path.join(DATA_JSON_FILES_PATH, "train.jsonl.gz"), mode="rt") as f:
        train_dataset = [example for line in f for example in prompt_formatter(json.loads(line))]
    with gzip.open(os.path.join(DATA_JSON_FILES_PATH, "test.jsonl.gz"), mode="rt") as f:
        test_dataset = [example for line in f for example in prompt_formatter(json.loads(line))]
    return train_dataset, test_dataset

def run_training(train_dataset, test_dataset, tokenizer, resume_from_checkpoint):
    base_model = transformers.AutoModelForCausalLM.from_pretrained(BASE_MODEL_PATH, load_in_8bit=True, low_cpu_mem_usage=True, use_safetensors=True)  # load in 8-bit quantized mode
    peft.prepare_model_for_kbit_training(base_model)
    if resume_from_checkpoint is not None:
        checkpoint_name = os.path.join(OUTPUT_DIRECTORY, resume_from_checkpoint)
        assert os.path.exists(checkpoint_name), checkpoint_name
        peft_base_model = peft.PeftModel.from_pretrained(base_model, checkpoint_name, is_trainable=True)  # TODO: is there a way to make this error out if it isn't in safetensors format?
    else:
        peft_base_model = peft.get_peft_model(base_model, peft.LoraConfig(
            # by default, this adds LoRA adapters around Llama2's "q_proj" and "v_proj", see https://github.com/huggingface/peft/blob/5a0e19dda1048ff8caaa12970ba7574f9cdfbf76/src/peft/utils/other.py#L280 for more details
            # some other models add more adapters, such as Guanaco, which does it on every linear layer except the head (https://github.com/artidoro/qlora/blob/845188de110d8eb7c95cc8907b54d8cb2e7c01bd/qlora.py#L221), but this doesn't seem to have too noticeable a benefit compared to models that just use these default settings, like Airoboros does (https://github.com/jondurbin/FastChat/blob/5bd738586ae6a495bd73152d74969465f30d43ac/fastchat/train/train_lora.py#L51)
            r=128,  # number of LoRA attention dimension parameters - directly proportional to LoRA adapter VRAM usage, set this to the largest value your hardware can support without running out of memory
            lora_alpha=16,  # alpha parameter for LoRA, essentially scales all of the LoRA weights, which determines "how much of an effect" this LoRA has on the final model
            lora_dropout=0.05,  # dropout probability for LoRA layers
            task_type=peft.TaskType.CAUSAL_LM,
        ))
    peft_base_model.print_trainable_parameters()
    torch.cuda.empty_cache()  # helps reduce VRAM usage - it's right after the PEFT version of the model was created, so some old stuff is still around unnecessarily

    trainer = trl.DPOTrainer(
        peft_base_model,
        base_model,  # use the original base model for comparison purposes (this model will be used in evaluation/inference mode, as a reference)
        args=transformers.TrainingArguments(
            output_dir=OUTPUT_DIRECTORY,
            dataloader_drop_last=True,  # when the dataset size isn't evenly divisible by the batch size, the remainder forms an incomplete batch - throw away this batch to avoid having to ever see an incomplete batch
            max_steps=TRAINING_STEPS,  # perform a fixed number of training steps before stopping
            evaluation_strategy="steps", eval_steps=200,  # run an evaluation every 200 training steps (~1 hour)
            save_strategy="steps", save_steps=200, save_safetensors=True,  # save a checkpoint every 200 training steps (~1 hour)
            logging_strategy="steps", logging_steps=1,  # log output every training step
            per_device_train_batch_size=BATCH_SIZE,
            per_device_eval_batch_size=BATCH_SIZE,
            learning_rate=1e-5, warmup_steps=150,  # linearly ramp up the learning rate for the AdamW optimizer from 0 to 1e-5 over the first 150 steps, then keep it at 1e-5 afterwards
            gradient_accumulation_steps=16,
            gradient_checkpointing=True,  # use gradient checkpointing to decrease VRAM usage
            bf16=True,  # use 16-bit bfloats for training instead of 32-bit floats in most operations (some are still kept in 32-bit for precision) to decrease VRAM usage and increase training performance, in practice the precision loss has a relatively small effect on the final result
            tf32=True,  # in newer NVIDIA hardware, this replaces the remaining 32-bit operations with a 19-bit TensorFloat operations to increase training performance, in practice the precision loss has no noticeable effect on the final result
            remove_unused_columns=False,  # the DPO default data collator requires this setting, since it customizes the model training/prediction steps
            run_name="llama2-direct-preference-finetuning",
            # TODO: when Apex becomes more stable + easier to install, look into using adamw_apex_fused rather than adamw_hf for the optim= parameter (note that some code uses optimizers= on the DPOTrainer itself, which overrides the optim= parameter on TrainingArguments, see https://github.com/huggingface/transformers/blob/53e1f5cf66d320b9c809f3940c707b6fef435d2d/src/transformers/trainer.py#L1084)
        ),
        beta=0.1,  # parameter controlling the deviation from the reference model, higher values prevent the model from deviating too far from the reference model, 0.1 is a relatively low value so the model will change quite a bit
        train_dataset=train_dataset,
        eval_dataset=test_dataset,
        tokenizer=tokenizer,
        max_length=CONTEXT_WINDOW_SIZE,  # any inputs that are longer than this value get truncated - first by cutting off the start of the prompt, then by cutting off the end of the response
        max_prompt_length=CONTEXT_WINDOW_SIZE // 2,  # when truncating by cutting off the start of the prompt, cut it down to half the context window size
    )

    os.makedirs(OUTPUT_DIRECTORY, exist_ok=True)
    trainer.train()
    peft_base_model.save_pretrained(os.path.join(OUTPUT_DIRECTORY, "final_checkpoint"), safe_serialization=True)  # save LoRA by itself

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--resume_from_checkpoint", type=str, default=None, help=f"If specified, start the training from the specified checkpoint (e.g., checkpoint-500). You can get a list of all checkpoints by running: ls {OUTPUT_DIRECTORY}'")
    args = parser.parse_args()

    transformers.set_seed(RANDOMNESS_SEED)

    tokenizer = transformers.AutoTokenizer.from_pretrained(BASE_MODEL_PATH, use_safetensors=True)
    tokenizer.pad_token = tokenizer.eos_token  # the padding token isn't set in the included tokenizer by default (see tokenizer.special_tokens_map for existing special tokens), set it manually
    train_dataset, test_dataset = create_datasets(tokenizer)
    run_training(train_dataset, test_dataset, tokenizer, args.resume_from_checkpoint)


import json
from PIL import Image
import cv2
import numpy as np

with open('data-5000.json') as f:
    data = json.load(f)
data = data.replace("\'", "\"")
data = json.loads(data)
keys = data.keys()

# read image file name & coordinates for cropping
for key in keys:
    img_path = 'D:\\github\\YOLOv4-darknet\\images\\' + key
    out_path = 'D:\\github\\YOLOv4-darknet\\cropped\\' + key
    img = cv2.imread(img_path)
    dim = img.shape
    idx = 0
    box_num = len(data[key])
    if box_num == 0:
        print(cv2.imwrite(out_path, img))
        continue
    elif box_num > 1:
        conf_arr = []
        for i in range(box_num):
            conf_arr.append(data[key][i]['confidence'])
        idx = [np.argmax(conf_arr)]
        idx = idx[0]
    x1 = float(data[key][idx]['x1'])
    x2 = float(data[key][idx]['x2'])
    y1 = float(data[key][idx]['y1'])
    y2 = float(data[key][idx]['y2'])
    if x1 < 0: x1 = 0
    if y1 < 0: y1 = 0
    if x2 > dim[1]: x2 = dim[1]
    if y2 > dim[0]: y2 = dim[0]
    img = img[int(y1):int(y2)+1, int(x1):int(x2)+1]
    print(cv2.imwrite(out_path, img))


from collections import deque

with open('input.txt') as f:
    data = f
    d = deque([0] * 9)
    for i in map(int, data.read().split(",")):
        d[i] += 1

for part in (80, 256 - 80):
    for _ in range(part):
        d.append(d.popleft())
        d[6] += d[-1]
    print(sum(d))


import os
import shutil
from pathlib import Path

# %%
def get_roi_data(data_path='/media/liuhy/data/数据集/ImageNet'):
    path_dst = Path.cwd().parent / 'data/ImageNet'
    path_label_file = Path(data_path) / 'labels'
    path_imgs = Path(data_path) / 'ILSVRC2012_img_train'

    tars = []
    with open(path_label_file) as f:
        line = f.readline()
        while line:
            tar_name = line.split(' ')[0] + '.tar'
            tars.append(tar_name)
            line = f.readline()

    for a_tar in tars:
        src = path_imgs / a_tar
        dst = path_dst / a_tar
        shutil.copyfile(src, dst)
    print('done copy')

get_roi_data()

# %%


import itertools
import operator
import os
import sys

import matplotlib as mpl
import numpy as np
import torch
from sklearn import metrics
from sklearn.metrics import classification_report
from torch.utils import data
from tqdm import tqdm

import config
import tensorboardwriter
from dataset.HisCancer_dataset import HisCancerDataset
from dataset.HisCancer_dataset import train_collate, val_collate
from gpu import gpu_profile
from loss.f import f1_macro, differenciable_f_softmax
from loss.focal import focalloss_softmax
from lr_scheduler.PlateauCyclicRestart import PlateauCyclicRestart
from project.HisCancer_project import HisCancer_net
from utils import load
from utils.load import save_checkpoint_fold, load_checkpoint_all_fold, save_onnx, remove_checkpoint_fold, set_milestone
from utils.lr_finder import LRFinder

if os.environ.get('DISPLAY', '') == '':
    print('WARNING: No display found. Using non-interactive Agg backend for loading matplotlib.')
    mpl.use('Agg')

class HisCancerValidation:
    def __init__(self, writer):
        self.writer = writer

        self.optimizers = []
        self.nets = []
        self.lr_schedulers = []
        self.train_loader = []
        self.validation_loader = []

        self.dataset = HisCancerDataset(config.DIRECTORY_CSV, config.DIRECTORY_SAMPLE_CSV, load_strategy="train", writer=self.writer, column='Target')
        self.folded_samplers = self.dataset.get_wsl_samples(fold=config.MODEL_FOLD)

        for fold in range(config.MODEL_FOLD):
            if fold not in config.train_fold:
                print("     Skipping dataset = HisCancerDataset(config.DIRECTORY_CSV, fold: #{})".format(fold))
                self.optimizers.append(None)
                self.nets.append(None)
                self.lr_schedulers.append(None)
            else:
                print("     Creating Fold: #{}".format(fold))
                net = HisCancer_net.se_resnext50_32x4d(config.TRAIN_NUM_CLASS, pretrained="imagenet", dropout_p=0.9)

                """FREEZING LAYER"""
                for i, c in enumerate(net.children()):
                    if len(config.MODEL_NO_GRAD) > i:
                        l = config.MODEL_NO_GRAD[i]
                        for child_counter, child in enumerate(list(c.children())):
                            if child_counter in l or l == [-1]:
                                print("Disable Gradient for child_counter: {}".format(child_counter))
                                for paras in child.parameters():
                                    paras.requires_grad = False
                            else:
                                print("Enable Gradient for child_counter: {}".format(child_counter))
                    else:
                        print("Enable Gradient for layer: {} (default)".format(i))

                if config.TRAIN_GPU_ARG: net = torch.nn.DataParallel(net, device_ids=config.TRAIN_GPU_LIST)
                optimizer = torch.optim.SGD(params=net.parameters(), lr=config.MODEL_INIT_LEARNING_RATE, momentum=config.MODEL_MOMENTUM, weight_decay=config.MODEL_WEIGHT_DECAY)
                self.optimizers.append(optimizer)
                self.nets.append(net)
                self.lr_schedulers.append(PlateauCyclicRestart(optimizer, eval_mode='max', factor=0.2, patience=1, verbose=False, threshold=1e-4, threshold_mode='abs', cooldown=0, eps=1e-8, base_lr=0.0023, max_lr=0.0069, step_size=2705, mode='plateau_cyclic', gamma=1., scale_mode='cycle', last_batch_iteration=-1, reduce_restart=3))

            self.train_loader.append(data.DataLoader(self.dataset,
                                           batch_size=config.MODEL_BATCH_SIZE,
                                           shuffle=False,
                                           sampler=self.folded_samplers[fold]["train"],
                                           batch_sampler=None,
                                           num_workers=config.TRAIN_NUM_WORKER,
                                           collate_fn=train_collate,
                                           pin_memory=True,
                                           drop_last=False,
                                           timeout=0,
                                           worker_init_fn=None,
                                           ))
            self.validation_loader.append(data.DataLoader(self.dataset,
                                                             batch_size=config.MODEL_BATCH_SIZE,
                                                             shuffle=False,
                                                             sampler=self.folded_samplers[fold]["val"],
                                                             batch_sampler=None,
                                                             num_workers=config.TRAIN_NUM_WORKER,
                                                             collate_fn=val_collate,
                                                             pin_memory=False,
                                                             drop_last=False,
                                                             timeout=0,
                                                             worker_init_fn=None,
                                                             ))
        if config.TRAIN_LOAD_OPTIMIZER: load_checkpoint_all_fold(self.nets, self.optimizers, self.lr_schedulers, config.DIRECTORY_LOAD)
        set_milestone(config.DIRECTORY_LOAD)

        """RESET LR"""
        if config.resetlr != 0:
            print("Reset Learning Rate to {}".format(config.resetlr))
            for optim in self.optimizers:
                if optim == None:
                    continue
                for g in optim.param_groups:
                    g['lr'] = config.resetlr

        if config.DISPLAY_SAVE_ONNX and config.DIRECTORY_LOAD: save_onnx(self.nets[config.train_fold[0]], (config.MODEL_BATCH_SIZE, 4, config.AUGMENTATION_RESIZE, config.AUGMENTATION_RESIZE), config.DIRECTORY_LOAD + ".onnx")

        if config.debug_lr_finder:
            val_loader = data.DataLoader(self.dataset,
                                         batch_size=config.MODEL_BATCH_SIZE,
                                         shuffle=False,
                                         sampler=self.folded_samplers[config.train_fold[0]]["val"],
                                         batch_sampler=None,
                                         num_workers=config.TRAIN_NUM_WORKER,
                                         collate_fn=val_collate,
                                         pin_memory=False,
                                         drop_last=False,
                                         timeout=0,
                                         worker_init_fn=None,
                                         ) if config.FIND_LR_ON_VALIDATION else None
            lr_finder = LRFinder(self.nets[config.train_fold[0]], self.optimizers[config.train_fold[0]], torch.nn.BCEWithLogitsLoss(), writer=self.writer, device="cuda")
            lr_finder.range_test(data.DataLoader(self.dataset,
                                                 batch_size=config.MODEL_BATCH_SIZE,
                                                 shuffle=False,
                                                 sampler=self.folded_samplers[config.train_fold[0]]["train"],
                                                 batch_sampler=None,
                                                 num_workers=config.TRAIN_NUM_WORKER,
                                                 collate_fn=train_collate,
                                                 pin_memory=True,
                                                 drop_last=False,
                                                 timeout=0,
                                                 worker_init_fn=None,
                                                 ), val_loader=val_loader, end_lr=0.1, num_iter=config.FIND_LR_RATIO, step_mode="exp")
            tensorboardwriter.write_plot(self.writer, lr_finder.plot(skip_start=0, skip_end=0, log_lr=False), "lr_finder-Linear")
            tensorboardwriter.write_plot(self.writer, lr_finder.plot(skip_start=0, skip_end=0, log_lr=True), "lr_finder-Log")
            lr_finder.reset()
            return

        """FREEZE DETECT"""
        for c in self.nets[config.train_fold[0]].children():
            for child_counter, child in enumerate(c.children()):
                req_grad = True
                for p in child.parameters():
                    if not p.requires_grad:
                        req_grad = False
                print("=======================Start Child Number #{} Grad: {}=======================".format(child_counter, req_grad))
                print("{}".format(child))
                print("=======================End Child Number #{} Grad: {}=======================".format(child_counter, req_grad))

        self.run()

    def run(self):
        try:
            for epoch in range(config.MODEL_EPOCHS):
                self.step_epoch(nets=self.nets,
                                optimizers=self.optimizers,
                                lr_schedulers=self.lr_schedulers,
                                batch_size=config.MODEL_BATCH_SIZE
                                )
                if config.TRAIN_GPU_ARG: torch.cuda.empty_cache()

            if config.TRAIN_GPU_ARG: torch.cuda.empty_cache()

        except KeyboardInterrupt as e:
            print(e)
            self.writer.close()
            print("To Resume: python train.py --versiontag 'test' --projecttag " + config.PROJECT_TAG + " --loadfile " + config.lastsave)
            try:
                sys.exit(0)
            except SystemExit:
                os._exit(0)

    def step_epoch(self,
                   nets,
                   optimizers,
                   lr_schedulers,
                   batch_size
                   ):
        config.epoch = config.epoch + 1

        evaluation = HisCancerEvaluation(self.writer, self.dataset.multilabel_binarizer)
        for fold, (net, optimizer, lr_scheduler) in enumerate(zip(nets, optimizers, lr_schedulers)):
            if net == None or optimizer == None or lr_scheduler == None:
                continue

            """UNFREEZE"""
            if config.epoch > config.MODEL_FREEZE_EPOCH:
                updated_children = []
                for i, c in enumerate(net.children()):
                    if len(config.MODEL_NO_GRAD) > i:
                        l = config.MODEL_NO_GRAD[i]
                        for child_counter, child in enumerate(list(c.children())):
                            if child_counter in l or l == [-1]:
                                for paras in child.parameters():
                                    if not paras.requires_grad:
                                        updated_children.append(child_counter)
                                        paras.requires_grad = True
                if len(updated_children) !=0:
                    print("Enable Gradient for child_counter: {}".format(updated_children))
                    tensorboardwriter.write_text(self.writer, "Unfreeze {} layers at epoch: {}".format(updated_children, config.epoch), config.global_steps[fold])
                # if config.MODEL_LEARNING_RATE_AFTER_UNFREEZE != 0:
                #     print("Reset Learning Rate to {}".format(config.MODEL_LEARNING_RATE_AFTER_UNFREEZE))
                #     for g in optimizer.param_groups:
                #         g['lr'] = config.MODEL_LEARNING_RATE_AFTER_UNFREEZE

            # import pdb; pdb.set_trace() #1357Mb -> 1215Mb
            """Switch Optimizers"""
            net = net.cuda()
            optimizer = load.move_optimizer_to_cuda(optimizer)
            if config.TRAIN_GPU_ARG: torch.cuda.empty_cache()
            val_loss, val_f1 = evaluation.eval_fold(net, self.validation_loader[config.fold])
            print("""
            ValidLoss: {}, ValidF1: {}
            """.format(val_loss, val_f1))
            net = net.cpu()
            optimizer = load.move_optimizer_to_cpu(optimizer) #3299Mb
            if config.TRAIN_GPU_ARG: torch.cuda.empty_cache() #1215Mb
        """LOSS"""
        f1 = f1_macro(evaluation.epoch_pred, evaluation.epoch_label).mean()
        f1_2 = metrics.f1_score((evaluation.epoch_label > config.EVAL_THRESHOLD).astype(np.byte), (evaluation.epoch_pred > config.EVAL_THRESHOLD).astype(np.byte), average='macro')  # sklearn does not automatically import matrics.
        f1_dict = dict(("Class-{}".format(i), x) for i, x in enumerate(metrics.f1_score((evaluation.epoch_label > config.EVAL_THRESHOLD).astype(np.byte), (evaluation.epoch_pred > config.EVAL_THRESHOLD).astype(np.byte), average=None)))
        f1_dict.update({"EvalF1": f1, "Sklearn": f1_2})

        # IT WILL MESS UP THE RANDOM SEED (CAREFUL)
        shakeup = dict()
        for i in range(100):
            public_lb = set(np.random.choice(range(int(len(evaluation.epoch_pred)*0.5)), int(len(evaluation.epoch_pred)*0.5), replace=False))
            private_lb = set(range(len(evaluation.epoch_pred)))-public_lb
            public_lb = list(public_lb)
            public_lb = metrics.roc_auc_score(evaluation.epoch_label[public_lb], evaluation.epoch_pred[public_lb])
            private_lb = list(private_lb)
            private_lb = metrics.roc_auc_score(evaluation.epoch_label[private_lb], evaluation.epoch_pred[private_lb])
            score_diff = private_lb-public_lb
            shakeup[score_diff] = (public_lb, private_lb)
        shakeup_keys = sorted(shakeup)
        shakeup_mean, shakeup_std = np.mean(shakeup_keys), np.std(shakeup_keys)
        tensorboardwriter.write_shakeup(self.writer, shakeup, shakeup_keys, config.epoch)

        soft_auc_macro = metrics.roc_auc_score(evaluation.epoch_label, evaluation.epoch_pred)
        hard_auc_macro = metrics.roc_auc_score((evaluation.epoch_label > config.EVAL_THRESHOLD).astype(np.byte), (evaluation.epoch_pred>config.EVAL_THRESHOLD).astype(np.byte))
        soft_auc_micro = metrics.roc_auc_score(evaluation.epoch_label, evaluation.epoch_pred, average='micro')
        hard_auc_micro = metrics.roc_auc_score((evaluation.epoch_label > config.EVAL_THRESHOLD).astype(np.byte), (evaluation.epoch_pred>config.EVAL_THRESHOLD).astype(np.byte), average='micro')

        report = classification_report(np.argmax(evaluation.epoch_label, axis=1), np.argmax(evaluation.epoch_pred, axis=1), target_names=["Negative", "Positive"])
        max_names = max(f1_dict.items(), key=operator.itemgetter(1))
        min_names = min(f1_dict.items(), key=operator.itemgetter(1))
        report = report + """
        Shakeup Mean: {}
        Shakeup STD: {}""".format(shakeup_mean, shakeup_std) + """
        Soft AUC Macro: {}
        Hard AUC Macro: {}
        Soft AUC Micro: {}
        Hard AUC Micro: {}
        """.format(soft_auc_macro, hard_auc_macro, soft_auc_micro, hard_auc_micro) + """
        F1 by sklearn = {}
        Max = {}, socre = {}
        Min = {}, score = {}
        """.format(f1_2, max_names[0], max_names[1], min_names[0], min_names[1])
        print(report)
        for lr_scheduler in lr_schedulers:
            if lr_scheduler == None:
                continue
            lr_scheduler.step(metrics=soft_auc_macro, epoch=config.epoch)
        tensorboardwriter.write_text(self.writer, report, config.epoch)

        tensorboardwriter.write_epoch_loss(self.writer, f1_dict, config.epoch)
        tensorboardwriter.write_pred_distribution(self.writer, evaluation.epoch_pred.flatten(), config.epoch)

        """THRESHOLD"""
        if config.EVAL_IF_THRESHOLD_TEST:
            best_threshold = 0.0
            best_val = 0.0

            best_threshold_dict = np.zeros(config.TRAIN_NUM_CLASS)
            best_val_dict = np.zeros(config.TRAIN_NUM_CLASS)

            pbar = tqdm(config.EVAL_TRY_THRESHOLD)
            for threshold in pbar:
                score = f1_macro(evaluation.epoch_pred, evaluation.epoch_label, thresh=threshold).mean()
                tensorboardwriter.write_threshold(self.writer, -1, score, threshold * 1000.0, config.fold)
                if score > best_val:
                    best_threshold = threshold
                    best_val = score
                pbar.set_description("Threshold: {}; F1: {}".format(threshold, score))

                for c in range(config.TRAIN_NUM_CLASS):
                    score = metrics.f1_score(evaluation.epoch_label[:][c], (evaluation.epoch_pred[:][c] > threshold))
                    tensorboardwriter.write_threshold(self.writer, c, score, threshold * 1000.0, config.fold)
                    if score > best_val_dict[c]:
                        best_threshold_dict[c] = threshold
                        best_val_dict[c] = score

            tensorboardwriter.write_best_threshold(self.writer, -1, best_val, best_threshold, config.epoch, config.fold)
            for c in range(config.TRAIN_NUM_CLASS): tensorboardwriter.write_best_threshold(self.writer, c, best_val_dict[c], best_threshold_dict[c], config.epoch, config.fold)

        """HISTOGRAM"""
        if config.DISPLAY_HISTOGRAM:
            tensorboardwriter.write_eval_loss(self.writer, {"EvalFocalMean": evaluation.mean(), "EvalFocalSTD": evaluation.std()}, config.epoch)
            tensorboardwriter.write_loss_distribution(self.writer, np.array(list(itertools.chain.from_iterable(evaluation.epoch_losses))).flatten(), config.epoch)

        """CLEAN UP"""
        del evaluation

class HisCancerEvaluation:
    def __init__(self, writer, binarlizer):
        self.writer = writer
        self.binarlizer = binarlizer
        """
        loss of one fold
        self.fold_losses = [
            eval_epoch -> [... losses of one batch...]
            eval_epoch -> [... losses of one batch...]
            eval_epoch -> [... losses of one batch...]
            eval_epoch -> [... losses of one batch...]
            eval_epoch -> [... losses of one batch...]
            eval_epoch -> [... losses of one batch...]
        ]
        """
        if config.DISPLAY_HISTOGRAM: self.epoch_losses = []  # [loss.flatten()]
        self.mean_losses = []
        # self.epoch_dict = np.array([]) # [fold_loss_dict]
        self.f1_losses = np.array([])

        # self.best_id = None
        # self.worst_id = None
        # self.best_loss = None
        # self.worst_loss = None

        self.epoch_pred = None
        self.epoch_label = None

    def eval_fold(self, net, validation_loader):
        focal_losses = np.array([])
        predict_total = None
        label_total = None

        # self.best_id = []
        # self.worst_id = []
        # self.best_loss = []
        # self.worst_loss = []

        print("Set Model Trainning mode to trainning=[{}]".format(net.eval().training))
        for eval_index in tqdm(range(config.EVAL_RATIO)):
            config.eval_index = eval_index
            pbar = tqdm(validation_loader)
            total_confidence = 0

            for batch_index, (ids, image, labels_0, image_for_display) in enumerate(pbar):

                """CALCULATE LOSS"""
                if config.TRAIN_GPU_ARG:
                    image = image.cuda()
                    labels_0 = labels_0.cuda()

                logits_predict = net(image)
                if config.TRAIN_GPU_ARG: torch.cuda.empty_cache()
                prob_predict = torch.nn.Softmax()(logits_predict)

                """LOSS"""
                focal = focalloss_softmax(alpha=0.25, gamma=5, eps=1e-7)(labels_0, logits_predict)
                f1, precise, recall = differenciable_f_softmax(beta=1)(labels_0, logits_predict)

                """EVALUATE LOSS"""
                focal = focal.detach()
                # focal_min = focal.min().item()
                # focal_min_id = (focal == focal_min).nonzero()
                # focal_min_id = focal_min_id.view(focal_min_id.size(), -1)[0]
                # focal_min_id = ids[focal_min_id.cpu().numpy()[0]]
                # focal_max = focal.max().item()
                # focal_max_id = (focal == focal_max).nonzero()
                # focal_max_id = focal_max_id.view(focal_max_id.size(), -1)[0]
                # focal_max_id = ids[focal_max_id.cpu().numpy()[0]]
                # self.best_loss = np.append(self.best_loss, focal_min)
                # self.worst_loss = np.append(self.worst_loss, focal_max)
                # self.best_id = np.append(self.best_id, focal_min_id)
                # self.worst_id = np.append(self.worst_id, focal_max_id)
                # del focal_min, focal_min_id, focal_max, focal_max_id
                if config.TRAIN_GPU_ARG: torch.cuda.empty_cache()

                """DETATCH"""
                focal = focal.cpu().numpy()
                focal_mean = focal.mean()
                f1 = f1.detach().cpu().numpy()
                precise = precise.detach().cpu().numpy().mean()
                recall = recall.detach().cpu().numpy().mean()
                # bce = bce.detach().cpu().numpy().mean()
                # positive_bce = positive_bce.detach().cpu().numpy().mean()
                # loss = loss.detach().cpu().numpy()
                labels_0 = labels_0.cpu().numpy()
                image = image.cpu().numpy()
                image_for_display = image_for_display.numpy()
                logits_predict = logits_predict.detach().cpu().numpy()
                prob_predict = prob_predict.detach().cpu().numpy()

                """SUM"""
                # np.append(self.f1_losses, f1_macro(prob_predict, labels_0).mean())
                self.f1_losses = np.append(self.f1_losses, f1.mean())
                focal_losses = np.append(focal_losses, focal_mean)

                confidence = np.absolute(prob_predict - 0.5).mean() + 0.5
                total_confidence = total_confidence + confidence

                """PRINT"""
                # label = np.array(self.dataset.multilabel_binarizer.inverse_transform(labels_0)[0])
                # pred = np.array(self.dataset.multilabel_binarizer.inverse_transform(prob_predict>0.5)[0])
                # pbar.set_description_str("(E{}-F{}) Stp:{} Label:{} Pred:{} Left:{}".format(int(config.global_steps[fold]), label, pred, left))
                pbar.set_description("(E{}F{}I{}) Focal:{} F1:{} Conf:{}".format(config.epoch, config.fold, config.eval_index, focal_mean, f1.mean(), total_confidence/(batch_index+1)))
                # if config.DISPLAY_HISTOGRAM: self.epoch_losses.append(focal.flatten())
                predict_total = np.concatenate((predict_total, prob_predict), axis=0) if predict_total is not None else prob_predict
                label_total = np.concatenate((label_total, labels_0), axis=0) if label_total is not None else labels_0

                """CLEAN UP"""
                del ids, image, image_for_display
                del focal, f1, precise, recall, labels_0, logits_predict, prob_predict
                if config.TRAIN_GPU_ARG: torch.cuda.empty_cache()
                if config.DEBUG_TRAISE_GPU: gpu_profile(frame=sys._getframe(), event='line', arg=None)
            del pbar
            if config.TRAIN_GPU_ARG: torch.cuda.empty_cache()
        """LOSS"""
        f1 = f1_macro(predict_total, label_total).mean()
        tensorboardwriter.write_eval_loss(self.writer, {"FoldFocal/{}".format(config.fold): focal_losses.mean(), "FoldF1/{}".format(config.fold): f1}, config.epoch)
        tensorboardwriter.write_pr_curve(self.writer, label_total, predict_total, config.epoch, config.fold)
        self.epoch_pred = np.concatenate((self.epoch_pred, predict_total), axis=0) if self.epoch_pred is not None else predict_total
        self.epoch_label = np.concatenate((self.epoch_label, label_total), axis=0) if self.epoch_label is not None else label_total

        if config.TRAIN_GPU_ARG: torch.cuda.empty_cache()
        mean_loss = focal_losses.mean()
        self.mean_losses.append(mean_loss)
        if config.TRAIN_GPU_ARG: torch.cuda.empty_cache()
        return mean_loss, f1

    def __int__(self):
        return self.mean()

    def mean(self, axis=None):
        # if axis == None: return np.array(list(itertools.chain.from_iterable(self.epoch_losses))).mean()
        # print("WARNING: self.epoch_losse may have different shape according to different shape of loss caused by different batch. Numpy cannot take the mean of it is baches shapes are different.")
        # return np.array(self.epoch_losses).mean(axis)
        return np.array(self.mean_losses).mean()

    def std(self, axis=None):
        if axis == None: return np.array(list(itertools.chain.from_iterable(self.epoch_losses))).std()
        print("WARNING: self.epoch_losse may have different shape according to different shape of loss caused by different batch. Numpy cannot take the mean of it is baches shapes are different.")
        return np.array(self.epoch_losses).std(axis)

    def f1_mean(self):
        return self.f1_losses.mean()


import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
import pickle

train = pd.read_csv('./datasets/traindata.csv')

X_train = train['comment_text']
Y_train = train[train.columns[4:]]

tfv = TfidfVectorizer(min_df=3, max_df=0.9, max_features=None,
                      strip_accents='unicode', analyzer='word', token_pattern=r'\w{1,}',
                      ngram_range=(1, 2), use_idf=1, smooth_idf=1,
                      sublinear_tf=1, stop_words='english')

tfv.fit(X_train)

#saving fit model
filename = "fit_model.sav"
pickle.dump(tfv,open(filename,'wb'))
loaded_fit = pickle.load(open('fit_model.sav','rb'))

X_train = loaded_fit.transform(X_train)

#saving seven model
for i in range(Y_train.shape[1]):
    feature = Y_train.columns[i]
    clf = LogisticRegression(C=4.0, solver='sag')
    clf.fit(X_train, Y_train.iloc[:, i])
    filename = 'model_{}.sav'.format(feature)
    pickle.dump(clf,open(filename,'wb'))


import logging

class Logging(object):
    def __init__(self, config, widget=None):
        logging.basicConfig(level=logging.INFO,
                        format='%(asctime)s %(levelname)-8s %(message)s',
                        datefmt='%d.%m %H:%M',
                        filename='run.log',
#                        encoding = "UTF-8",
                        filemode='w')
        self.log = logging
        if not config:
            self.log.critical('! Cannot find config file !')
            raise
        self.display_critical = config['display_critical']
        self.display_debug = config['display_debug']
        self.display_info = config['display_info']
        self.display_warning = config['display_warning']

        self.messages = []
        self.widget = widget

    def set_widget(self, widget):
        self.widget = widget

    def indicate(self):
        disp_html = ''
        if len(self.messages) > 100:
            self.messages = self.messages[-100:]
        for mes in self.messages:
            if mes[0] == 'critical':
                disp_html += '
' + mes[1] + '
'
            elif mes[0] == 'warning':
                disp_html += '
' + mes[1] + '
'
            elif mes[0] == 'info':
                disp_html += '
' + mes[1] + '
'
            elif mes[0] == 'debug':
                disp_html += '
' + mes[1] + '
'
            else:
                raise
        disp_html = disp_html.replace(u'%margin%', u'margin:0; -qt-block-indent:0; text-indent:0')
        self.widget.setHtml(disp_html)
        try:
            scrollBar = self.widget.verticalScrollBar()
            scrollBar.setValue(scrollBar.maximum())
        except:
            print(Exception)
            pass

    def __print2lambdas(self, msg, if_check, if_force_display, log_lambda, type_enum):
        if  if_check or if_force_display:
            self.messages.append((type_enum, msg))
            if self.widget:
                self.indicate()
            else:
                print('{}: {}'.format(type_enum, msg))
        log_lambda(msg)

    def critical(self, msg, force_display=False):
        self.__print2lambdas(msg, self.display_critical, force_display, self.log.critical, 'critical')

    def debug(self, msg, force_display=False):
        self.__print2lambdas(msg, self.display_debug, force_display, self.log.debug, 'debug')

    def warning(self, msg, force_display=False):
        self.__print2lambdas(msg, self.display_warning, force_display, self.log.warning, 'warning')

    def info(self, msg, force_display=False):
        self.__print2lambdas(msg, self.display_info, force_display, self.log.info, 'info')

class Logging2Console(Logging):
    def __init__(self):
        pass

    def critical(self, msg, force_display=False):
        print('{}: {}'.format('critical', msg))

    def debug(self, msg, force_display=False):
        print('{}: {}'.format('debug', msg))

    def warning(self, msg, force_display=False):
        print('{}: {}'.format('warning', msg))

    def info(self, msg, force_display=False):
        print('{}: {}'.format('info', msg))


import concurrent.futures
import multiprocessing
import random
from abc import abstractmethod
from math import ceil, floor
from typing import Dict, List, Optional, Tuple, Union
from uu import Error

import torch
from rich.progress import Console, track
from torch.utils.data import Dataset
from torch.utils.data.dataloader import DataLoader

from nerfstudio.cameras.cameras import Cameras
from nerfstudio.cameras.rays import RayBundle
from nerfstudio.data.datasets.base_dataset import InputDataset
from nerfstudio.data.datasets.dynamic_dataset import DynamicDataset
from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate
from nerfstudio.utils.misc import get_dict_to_torch

CONSOLE = Console(width=120)

class CacheDataloader(DataLoader):
    """Collated image dataset that implements caching of default-pytorch-collatable data.
    Creates batches of the InputDataset return type.

    Args:
        dataset: Dataset to sample from.
        num_samples_to_collate: How many images to sample rays for each batch. -1 for all images.
        num_times_to_repeat_images: How often to collate new images. -1 to never pick new images.
        device: Device to perform computation.
        collate_fn: The function we will use to collate our training data
    """

    def __init__(
        self,
        dataset: Dataset,
        num_images_to_sample_from: int = -1,
        num_times_to_repeat_images: int = -1,
        device: Union[torch.device, str] = "cpu",
        collate_fn=nerfstudio_collate,
        **kwargs,
    ):
        self.dataset = dataset
        super().__init__(dataset=dataset, **kwargs)  # This will set self.dataset
        self.num_times_to_repeat_images = num_times_to_repeat_images
        self.cache_all_images = (num_images_to_sample_from == -1) or (num_images_to_sample_from >= len(self.dataset))
        self.num_images_to_sample_from = len(self.dataset) if self.cache_all_images else num_images_to_sample_from
        self.device = device
        self.collate_fn = collate_fn
        self.num_workers = kwargs.get("num_workers", 0)

        self.num_repeated = self.num_times_to_repeat_images  # starting value
        self.first_time = True
        self.iter_step = 0

        self.cached_collated_batch = None
        if self.cache_all_images:
            CONSOLE.print(f"Caching all {len(self.dataset)} images.")
            if len(self.dataset) > 500:
                CONSOLE.print(
                    "[bold yellow]Warning: If you run out of memory, try reducing the number of images to sample from."
                )
            self.cached_collated_batch = self._get_collated_batch()
            # Pre-compute importance sampling weights
            if isinstance(self.dataset, DynamicDataset) and self.dataset.use_importance_sampling:
                ist_w = None
                # IST weights in evaluation can be useful to compare how well the model does
                # on dynamic content and is more representative than
                ist_w = self.dataset.compute_is(self.cached_collated_batch, self.device, offline=True)
                self.cached_collated_batch["ist_weights"] = ist_w
        elif self.num_times_to_repeat_images == -1:
            CONSOLE.print(
                f"Caching {self.num_images_to_sample_from} out of {len(self.dataset)} images, without resampling."
            )
        else:
            CONSOLE.print(
                f"Caching {self.num_images_to_sample_from} out of {len(self.dataset)} images, "
                f"resampling every {self.num_times_to_repeat_images} iters."
            )

    def __getitem__(self, idx):
        return self.dataset.__getitem__(idx)

    def _get_batch_list(self):
        """Returns a list of batches from the dataset attribute."""

        # Modes:
        #  1. "normal": random samples
        #  2. "randsteps": random steps, same for each camera
        #  3. "lowfps": Lower FPS by k.
        pick_mode = "normal"
        if isinstance(self.dataset, DynamicDataset):
            pick_mode = self.dataset.pick_mode

        total_size = len(self.dataset)
        to_sample = self.num_images_to_sample_from

        if total_size == to_sample:
            pick_mode = "normal"

        indices: List

        if pick_mode == "normal":
            indices = random.sample(range(total_size), k=to_sample)
        elif pick_mode == "randsteps" or pick_mode == "lowfps":
            times = list(set(self.dataset.cameras.times.squeeze().tolist()))

            picked_times: List
            if pick_mode == "randsteps":
                cam_ids = self.dataset.cameras.ids.squeeze().tolist()

                # Unique cameras:
                nb_unique_cams = len(set(cam_ids))
                steps_to_pick = int(to_sample / nb_unique_cams)  # floor
                # Always include start and end
                picked_times = [times[0], times[-1]]
                picked_times += random.sample(times[1:-1], k=(steps_to_pick-2))
            elif pick_mode == "lowfps":
                # to_sample should be at the desired framerate
                if total_size % to_sample != 0:
                    print(
                        "[Warning | lowfps]: Ratio of images to sample over total is not an integer. Rounding down. This may result in a lower framerate than desired."
                    )

                # Pick every k frames
                k = ceil(total_size / to_sample)
                picked_times = times[::k]
                if len(times) % k != 0:
                    print(
                        "[Warning | lowfps]: Duration not dividable by desired framerate ratio. Rounding down. This may result in a lower framerate than desired."
                    )
                    # Remove one element
                    picked_times = picked_times[:-1]

                if total_size % to_sample != 0 or len(times) % k != 0:
                    print(
                        f"[Warning | lowfps]: effective FPS downscale chosen: {len(times) / len(picked_times)}. Desired was about {(total_size / to_sample):.4f}."
                    )

            indices = []
            for idx in range(total_size):
                if self.dataset.cameras.times[idx].item() in picked_times:
                    indices.append(idx)

            # Left to pick
            left_to_pick = to_sample - len(indices)
            if left_to_pick > 0:
                # print(f"[INFO]: Picking {left_to_pick} more images ({pick_mode})")
                range_wo_picked = [x for x in range(total_size) if x not in indices]
                indices += random.sample(range_wo_picked, k=left_to_pick)
        else:
            raise ValueError("Unknown pick_mode: " + pick_mode)

        # if isinstance(self.dataset, DynamicDataset):
        #    print(f"Picked {len(indices)} images with mode {pick_mode}.")

        if len(indices) != to_sample:
            raise Error("Not enough images to sample from.")

        batch_list = []
        results = []

        num_threads = int(self.num_workers) * 4
        num_threads = min(num_threads, multiprocessing.cpu_count() - 1)
        num_threads = max(num_threads, 1)

        with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
            for idx in indices:
                res = executor.submit(self.dataset.__getitem__, idx)
                results.append(res)

            for res in track(results, description="Loading data batch", transient=True):
                batch_list.append(res.result())

        return batch_list

    def _get_collated_batch(self):
        """Returns a collated batch."""
        batch_list = self._get_batch_list()
        collated_batch = self.collate_fn(batch_list)
        # TODO exclude depth image
        collated_batch = get_dict_to_torch(collated_batch, device=self.device, exclude=["image", "depth_image"])
        return collated_batch

    def __iter__(self):
        while True:
            if self.cache_all_images:
                collated_batch = self.cached_collated_batch
            elif self.first_time or (
                self.num_times_to_repeat_images != -1 and self.num_repeated >= self.num_times_to_repeat_images
            ):
                # trigger a reset
                self.num_repeated = 0
                collated_batch = self._get_collated_batch()

                # Pre-compute importance sampling weights
                if isinstance(self.dataset, DynamicDataset) and self.dataset.use_importance_sampling:
                    if self.iter_step + self.num_times_to_repeat_images >= self.dataset.iters_to_start_ist:
                        # Also perform on evaluation dataset (was not the case beforehand)
                        ist_w = self.dataset.compute_is(collated_batch, self.device)
                        collated_batch["ist_weights"] = ist_w

                # possibly save a cached item
                self.cached_collated_batch = collated_batch if self.num_times_to_repeat_images != 0 else None
                self.first_time = False
            else:
                collated_batch = self.cached_collated_batch
                self.num_repeated += 1
            self.iter_step += 1
            collated_batch["iter_steps"] = self.iter_step
            yield collated_batch

class EvalDataloader(DataLoader):
    """Evaluation dataloader base class

    Args:
        input_dataset: InputDataset to load data from
        device: Device to load data to
    """

    def __init__(
        self,
        input_dataset: InputDataset,
        device: Union[torch.device, str] = "cpu",
        **kwargs,
    ):
        self.input_dataset = input_dataset
        self.cameras = input_dataset.cameras.to(device)
        self.device = device
        self.kwargs = kwargs
        super().__init__(dataset=input_dataset)

    @abstractmethod
    def __iter__(self):
        """Iterates over the dataset"""
        return self

    @abstractmethod
    def __next__(self) -> Tuple[RayBundle, Dict]:
        """Returns the next batch of data"""

    def get_camera(self, image_idx: int = 0) -> Cameras:
        """Get camera for the given image index

        Args:
            image_idx: Camera image index
        """
        return self.cameras[image_idx]

    def get_data_from_image_idx(self, image_idx: int) -> Tuple[RayBundle, Dict]:
        """Returns the data for a specific image index.

        Args:
            image_idx: Camera image index
        """
        ray_bundle = self.cameras.generate_rays(camera_indices=image_idx, keep_shape=True)
        batch = self.input_dataset[image_idx]
        batch = get_dict_to_torch(batch, device=self.device, exclude=["image"])
        return ray_bundle, batch

class FixedIndicesEvalDataloader(EvalDataloader):
    """Dataloader that returns a fixed set of indices.

    Args:
        input_dataset: InputDataset to load data from
        image_indices: List of image indices to load data from. If None, then use all images.
        device: Device to load data to
    """

    def __init__(
        self,
        input_dataset: InputDataset,
        image_indices: Optional[Tuple[int]] = None,
        device: Union[torch.device, str] = "cpu",
        **kwargs,
    ):
        super().__init__(input_dataset, device, **kwargs)
        if image_indices is None:
            self.image_indices = list(range(len(input_dataset)))
        else:
            self.image_indices = image_indices
        self.count = 0

    def __iter__(self):
        self.count = 0
        return self

    def __next__(self):
        if self.count < len(self.image_indices):
            image_idx = self.image_indices[self.count]
            ray_bundle, batch = self.get_data_from_image_idx(image_idx)
            self.count += 1
            return ray_bundle, batch
        raise StopIteration

class RandIndicesEvalDataloader(EvalDataloader):
    """Dataloader that returns random images.

    Args:
        input_dataset: InputDataset to load data from
        device: Device to load data to
    """

    def __init__(
        self,
        input_dataset: InputDataset,
        device: Union[torch.device, str] = "cpu",
        **kwargs,
    ):
        super().__init__(input_dataset, device, **kwargs)
        self.count = 0

    def __iter__(self):
        self.count = 0
        return self

    def __next__(self):
        if self.count < 1:
            image_indices = range(self.cameras.size)
            image_idx = random.choice(image_indices)
            ray_bundle, batch = self.get_data_from_image_idx(image_idx)
            self.count += 1
            return ray_bundle, batch
        raise StopIteration


import RPi.GPIO as GPIO # Import Raspberry Pi GPIO library
import time
import led
import cv2
import requests
from tinydb import TinyDB, Query

#GPIO.setwarnings(False) # Ignore warning for now
GPIO.setmode(GPIO.BOARD) # Use physical pin numbering
GPIO.setup(10, GPIO.IN, pull_up_down=GPIO.PUD_DOWN) # Set pin 10 to be an input pin and set initial value to be pulled low (off)

class webcam:
    def __init__(self, path="./images/"):
        self.path = path
        pass

    def capture(self, num_images = 10, base_filename="cam"):
        camera = cv2.VideoCapture(0)

        for i in range(num_images):
            return_value, image = camera.read()
            filename = self.path + base_filename+'_'+str(i)+'.png'
            print(filename)
            cv2.imwrite(filename, image)
        del(camera)

class button:
    states = ['ready','started','stopped']
    button_state = 0
    button_state_time = 0
    button_delay = 1

    def __init__(self):
        if self.button_state_time == 0:
            self.button_state_time = time.time()

    def button_callback(self, channel):
        now = time.time()
        if ((now-self.button_state_time)>self.button_delay):
            print("Button was pushed!")
            self.button_state_time = now
            if self.button_state == 0:
                self.button_state = 1
            elif self.button_state == 1:
                self.button_state = 2
            elif self.button_state == 2:
                self.button_state = 0

    def state(self):
        return self.states[self.button_state]

def main():
    led.clear()

    db = TinyDB('./scores.json')

    print("Press button to start game\n\n")
    b = button()
    GPIO.add_event_detect(10,GPIO.RISING, callback=b.button_callback) # Setup event on pin 10 rising edge

    ready_state = False
    stopped_state = True
    while True:
        while b.state() =="ready":
            if not ready_state:
                led.display("00.0000")

                ready_state = True
                stopped_state = False
            start = time.time()
        while b.state() == "started":
            time.sleep(.1)
            end = time.time()
            elapsed = end - start
            float_elapsed = float(elapsed)
            (left,right)= str(float_elapsed).split(".")
            left = left.rjust(2,"0")
            right= right[:4]
            formatStr = left + "." + right
            led.display(formatStr)
        while b.state() == "stopped":
            if not stopped_state:
                timestr = time.gmtime() #strftime("%Y%m%d-%H%M%S")
                score = abs(float(formatStr)-10)
                obj = {'timestamp': timestr, 'time': formatStr, 'score':score}
                try:
                    r = requests.post('http://192.168.120.92:3000/game-end', data = obj)
                except:
                    print("dashboard down?")
                cam = webcam()
                cam.capture(base_filename=timestr)
                db.insert(obj)
                print("image captured: " + formatStr)
                stopped_state = True
                ready_state = False

        pass

    GPIO.cleanup() # Clean up

if __name__ == "__main__":
    main()


import telebot
import config
from telebot import types
import json

bot = telebot.TeleBot(config.TOKEN)
scr_data = {}

with open('data/pizza.json', 'rb') as file:
    data = json.load(file)
    scr_data['pizza'] = data

with open('data/salads.json', 'rb') as file:
    data = json.load(file)
    scr_data['salads'] = data

with open('data/desserts.json', 'rb') as file:
    data = json.load(file)
    scr_data['desserts'] = data

with open('data/beverages.json', 'rb') as file:
    data = json.load(file)
    scr_data['beverages'] = data

@bot.message_handler(commands=['start'])
def greeting(message):
    usr = message.from_user.to_dict()
    greeting = f'Привіт, {usr["first_name"]}!\nМене звуть Орися, я допоможу тобі зробити у нас замовлення!'

    markup = types.ReplyKeyboardMarkup(resize_keyboard=True)
    discounts = types.KeyboardButton('Акції')
    pizza = types.KeyboardButton('Піца')
    salads = types.KeyboardButton("Салати")
    dessert = types.KeyboardButton("Десерти")
    beverages = types.KeyboardButton("Напої")

    markup.add(discounts, pizza, salads, dessert, beverages)

    with open('sticker.webp', 'rb') as sti:
        bot.send_sticker(message.chat.id, sti)
        bot.send_message(message.chat.id, greeting, reply_markup=markup)

@bot.message_handler(content_types=['text'])
def menu(message):
    if message.chat.type == 'private':
        if message.text == 'Акції':
            discounts(message)
        elif message.text == 'Піца':
            pizza(message)
        elif message.text == 'Салати':
            salads(message)
        elif message.text == 'Десерти':
            desserts(message)
        elif message.text == 'Напої':
            beverages(message)

@bot.message_handler(content_types=['text'])
def discounts(message):
    with open('data/discount.json', 'rb') as file:
        data = json.load(file)
        for i in data:
            markup = types.InlineKeyboardMarkup(row_width=1)
            item = types.InlineKeyboardButton('Cкористатися', url=i['link'])

            markup.add(item)
            bot.send_photo(message.chat.id, i['img'])
            bot.send_message(message.chat.id, i['text'], reply_markup=markup)

@bot.message_handler(content_types=['text'])
def pizza(message):
    reply_text = 'Фотографії продуктів несуть ознайомчий характер і можуть відрізнятися від оригіналу'
    with open('data/pizza.json', 'rb') as file:
        data = json.load(file)

        markup = types.InlineKeyboardMarkup(row_width=4)
        for i in range(len(scr_data['pizza'])):
            btn_txt = f"{data[i]['name']} - {data[i]['price']}"
            item = types.InlineKeyboardButton(btn_txt, callback_data=str(scr_data['pizza'][i]['id']))
            markup.add(item)

        bot.send_photo(message.chat.id, data[0]['img'])
        bot.send_message(message.chat.id, text=reply_text, reply_markup=markup)

@bot.message_handler(content_types=['text'])
def salads(message):
    reply_text = 'Боїтеся помилитися та обрати не той салат? Зателефонуйте до доставки піци «Lа П’єц» у Львові та ' \
                 'проконсультуйтеся перед тим, як зробити замовлення. '

    with open('data/salads.json', 'rb') as file:
        data = json.load(file)

        markup = types.InlineKeyboardMarkup(row_width=4)
        for i in range(len(scr_data['salads'])):
            btn_txt = f"{data[i]['name']} - {data[i]['price']}"
            item = types.InlineKeyboardButton(btn_txt, callback_data=str(scr_data['salads'][i]['id']))
            markup.add(item)

        bot.send_photo(message.chat.id, data[0]['img'])
        bot.send_message(message.chat.id, text=reply_text, reply_markup=markup)

@bot.message_handler(content_types=['text'])
def desserts(message):
    reply_text = 'Десерти:'
    with open('data/desserts.json', 'rb') as file:
        data = json.load(file)

        markup = types.InlineKeyboardMarkup(row_width=4)
        for i in range(len(scr_data['desserts'])):
            btn_txt = f"{scr_data['desserts'][i]['name']} - {scr_data['desserts'][i]['price']}"
            item = types.InlineKeyboardButton(btn_txt, callback_data=scr_data['desserts'][i]['id'])
            markup.add(item)

        bot.send_photo(message.chat.id, data[0]['img'])
        bot.send_message(message.chat.id, text=reply_text, reply_markup=markup)

@bot.message_handler(content_types=['text'])
def beverages(message):
    reply_text = 'Напої:'
    with open('data/beverages.json', 'rb') as file:
        data = json.load(file)

        markup = types.InlineKeyboardMarkup(row_width=4)
        for i in range(len(scr_data['beverages'])):
            btn_txt = f"{scr_data['beverages'][i]['name']} - {scr_data['beverages'][i]['price']}"
            item = types.InlineKeyboardButton(btn_txt, callback_data=scr_data['beverages'][i]['id'])
            markup.add(item)

        bot.send_photo(message.chat.id, data[0]['img'])
        bot.send_message(message.chat.id, text=reply_text, reply_markup=markup)

@bot.callback_query_handler(func=lambda call: True)
def msg_callback(call):

    try:
        if call.message:
            for i in range(len(scr_data['pizza'])):
                if call.data == str(scr_data['pizza'][i]['id']):
                    msg_template = f'{scr_data["pizza"][i]["name"]}\n{scr_data["pizza"][i]["desc"]}\nЦіна: {scr_data["pizza"][i]["price"]}\nРозмір: {scr_data["pizza"][i]["size"]}\nВага: {scr_data["pizza"][i]["weight"]}'
                    markup = types.InlineKeyboardMarkup(row_width=4)
                    item = types.InlineKeyboardButton('Дивитись на сайті', url=scr_data['pizza'][i]['link'])
                    markup.add(item)
                    bot.send_photo(call.message.chat.id, scr_data['pizza'][i]['img'])
                    bot.send_message(call.message.chat.id, msg_template, reply_markup=markup)

            for i in range(len(scr_data['salads'])):
                if call.data == str(scr_data['salads'][i]['id']):
                    msg_template = f'{scr_data["salads"][i]["name"]}\n{scr_data["salads"][i]["desc"]}\nЦіна: {scr_data["salads"][i]["price"]}\nВага: {scr_data["salads"][i]["weight"]}'

                    bot.send_photo(call.message.chat.id, scr_data['salads'][i]['img'])
                    bot.send_message(call.message.chat.id, msg_template)
            for i in range(len(scr_data['beverages'])):
                if call.data == str(scr_data['beverages'][i]['id']):
                    msg_template = f'{scr_data["beverages"][i]["name"]}\nЦіна: {scr_data["beverages"][i]["price"]}'
                    # markup = types.InlineKeyboardMarkup(row_width=4)
                    # item = types.InlineKeyboardButton('Дивитись на сайті', url=scr_data['beverages'][i]['link'])
                    # markup.add(item)
                    bot.send_photo(call.message.chat.id, scr_data['beverages'][i]['img'])
                    bot.send_message(call.message.chat.id, msg_template)
            for i in range(len(scr_data['desserts'])):
                if call.data == str(scr_data['desserts'][i]['id']):
                    msg_template = f'{scr_data["desserts"][i]["name"]}\n{scr_data["desserts"][i]["desc"]}\nВага: {scr_data["desserts"][i]["desc"]}\nЦіна: {scr_data["desserts"][i]["price"]}'
                    # markup = types.InlineKeyboardMarkup(row_width=4)
                    # item = types.InlineKeyboardButton('Дивитись на сайті', url=scr_data['desserts'][i]['link'])
                    # markup.add(item)
                    bot.send_photo(call.message.chat.id, scr_data['desserts'][i]['img'])
                    bot.send_message(call.message.chat.id, msg_template)
    except Exception as e:
        print(e)

if __name__ == '__main__':
    bot.polling(none_stop=True)


import Tkinter
import ttk
import tkMessageBox
import threading
import logging
import re
import time

import thread_safe_tk_widgets
import fast_test
import slow_test

FAST_LOCKIN = "GPIB0::9::INSTR"
SLOW_LOCKIN = "GPIB0::11::INSTR"
DMM = "GPIB0::16::INSTR"

class fast_frame(Tkinter.Frame):
    def __init__(self, parent, fast_sample_Variable):
        Tkinter.Frame.__init__(self, parent, bd=10)
        self.parent = parent
        self.parent_fast_sample_Variable = fast_sample_Variable
        self.initialize()

    def initialize(self):
        # GUI part
        self.grid()

        captain_line = Tkinter.Label(self, text='Fast Measurement', font=("Helvetica", 16), anchor="w")
        captain_line.grid(column=0,row=0,columnspan=10)

        self.output_box = thread_safe_tk_widgets.ThreadSafeText(self, undo = False, wrap='none',
                          width = 60, height = 34)
        self.output_box.grid(column=0,row=1,columnspan=10,sticky='EW')

        initial_interval_Label = Tkinter.Label(self, text='Initial interval (sec):',anchor='e')
        initial_interval_Label.grid(column=0,row=5)
        self.initial_interval_Variable = Tkinter.StringVar()
        self.interval_select_box = ttk.Combobox(self, textvariable=self.initial_interval_Variable, state='readonly', width=10)
        self.interval_select_box.grid(column=1,row=5,sticky='EW')
        self.interval_select_box['values'] = ['0.01','0.05','0.1','0.3','1','3']
        self.initial_interval_Variable.set('0.01')

        sample_select_Label = Tkinter.Label(self, text='Fast mode sample:',anchor='e')
        sample_select_Label.grid(column=2,row=5)
        self.sample_select_Variable = Tkinter.StringVar()
        self.sample_select_box = ttk.Combobox(self, textvariable=self.sample_select_Variable, state='readonly', width=6)
        self.sample_select_box.grid(column=3,row=5,sticky='EW')
        self.sample_select_box['values'] = [str(i) for i in range(17)]
        self.sample_select_Variable.set('0')

        freq_entry_Label = Tkinter.Label(self, text='Initial freq:',anchor='e')
        freq_entry_Label.grid(column=4,row=5)
        vcmd = (self.register(self.EntryValidate),'%P')
        self.freq_Entry = Tkinter.Entry(self, validate="key", validatecommand=vcmd, width=10,state = Tkinter.NORMAL)
        self.freq_Entry.grid(column=5,row=5)
        self.freq_Entry.delete(0,Tkinter.END)
        self.freq_Entry.insert(0,'43')

        self.auto_tc_Var = Tkinter.IntVar()
        self.auto_tc_Var.set(0)
        auto_tc_button = Tkinter.Checkbutton(self, text='Auto interval and frequency', variable=self.auto_tc_Var)
        auto_tc_button.grid(column=0,row=6,columnspan=2,sticky='W')

        self.last_point_Var = Tkinter.IntVar()
        self.last_point_Var.set(0)
        last_point_button = Tkinter.Checkbutton(self, text='Last Point', variable=self.last_point_Var)
        last_point_button.grid(column=2,row=6,sticky='W')

        self.start_button_Variable = Tkinter.StringVar()
        self.start_button = ttk.Button(self,textvariable=self.start_button_Variable,command=self.OnButtonClick)
        self.start_button.grid(column=9,row=5,rowspan=2,sticky='EWNS')
        self.start_button_Variable.set('START')

        self.status_bar = thread_safe_tk_widgets.ThreadSafeLabel(self,text="Hello!",anchor="w",fg="white",bg="blue")
        self.status_bar.grid(column=0,row=7,columnspan=10,sticky='EW')

        # logic part
        self.test = fast_test.fast_test(0, FAST_LOCKIN, print_ch = 'Tk', init_freq = float(self.freq_Entry.get()),
                                        Tk_output = self.output_box, Tk_status = self.status_bar)
        self.auto_tc_Var.trace('w',self.OnAutoTcChange)
        self.last_point_Var.trace('w',self.OnLastPointChange)
        self.is_running = False

    def OnButtonClick(self):
        self.start_button.state(['disabled'])
        if not self.is_running: # not running, then start
            self.start_button_Variable.set('STOP') #1
            self.sample_select_box.state(['disabled']) #2
            self.interval_select_box.state(['disabled']) #3
            self.freq_Entry.configure(state='disabled') #7
            self.is_running = True #4
            self.test._to_stop = False #5
            self.test.sample_no = int(self.sample_select_Variable.get())
            self.parent_fast_sample_Variable.set(self.sample_select_Variable.get()) #6
            self.test.interval = float(self.initial_interval_Variable.get())
            self.test.freq = float(self.freq_Entry.get())
            #self.test.initialize()
            self.test_thread = threading.Thread(target = self.test.main_test_loop)
            self.test_thread.start()
        else: # already running, then stop
            self.test._to_stop = True #5
            self.test_thread.join()
            self.freq_Entry.configure(state='normal') #7
            self.parent_fast_sample_Variable.set('0') #6
            self.is_running = False #4
            self.interval_select_box.state(['!disabled']) #3
            self.sample_select_box.state(['!disabled']) #2
            self.start_button_Variable.set('START') #1
        self.start_button.state(['!disabled'])

    def OnAutoTcChange(self,*args):
        self.test._auto_tc = bool(self.auto_tc_Var.get())

    def OnLastPointChange(self,*args):
        self.test._last_point = bool(self.last_point_Var.get())

    def EntryValidate(self, P):
        # Disallow anything but float number
        valid = bool(re.match(r'^\d\.?\d*$',P))
        if not valid:
            self.bell()
        return valid

class slow_frame(Tkinter.Frame):
    def __init__(self, parent, slow_freq_Variable):
        Tkinter.Frame.__init__(self, parent, bd=10)
        self.parent = parent
        self.slow_freq_Variable = slow_freq_Variable
        self.initialize()

    def initialize(self):
        # GUI part
        self.grid()

        captain_line = Tkinter.Label(self, text='Slow Measurement', font=("Helvetica", 16))
        captain_line.grid(column=0,row=0,columnspan=10)

        self.output_box = thread_safe_tk_widgets.ThreadSafeText(self, undo = False, wrap='none',
                          width = 60, height = 31)
        self.output_box.grid(column=0,row=1,columnspan=10,sticky='EW')

        initial_interval_Label = Tkinter.Label(self, text='Initial interval (sec):',anchor='e')
        initial_interval_Label.grid(column=0,row=4)
        self.initial_interval_Variable = Tkinter.StringVar()
        self.interval_select_box = ttk.Combobox(self, textvariable=self.initial_interval_Variable, state='readonly', width=10)
        self.interval_select_box.grid(column=1,row=4,sticky='EW')
        self.interval_select_box['values'] = ['15','30','60','150','300','600']
        self.initial_interval_Variable.set('30')

        wait_time_Label = Tkinter.Label(self, text='Stabilize time (sec):',anchor='e')
        wait_time_Label.grid(column=2,row=4)
        self.wait_time_Variable = Tkinter.StringVar()
        self.wait_time_select_box = ttk.Combobox(self, textvariable=self.wait_time_Variable, state='readonly', width=6)
        self.wait_time_select_box.grid(column=3,row=4,sticky='EW')
        self.wait_time_select_box['values'] = ['1','3','5','10','30']
        self.wait_time_Variable.set('3')

        freq_entry_Label = Tkinter.Label(self, text='Slow freq:',anchor='e')
        freq_entry_Label.grid(column=4,row=4)
        vcmd = (self.register(self.EntryValidate),'%P')
        self.freq_Entry = Tkinter.Entry(self, textvariable=self.slow_freq_Variable, validate="key", validatecommand=vcmd, width=10)
        self.freq_Entry.grid(column=5,row=4)
        self.freq_Entry.delete(0,Tkinter.END)
        self.freq_Entry.insert(0,'1.3')

        self.sample_select_box = Tkinter.LabelFrame(self, text='Slow mode samples')
        self.sample_select_box.grid(column=0,row=5,columnspan=3,sticky='W')

        self.selected_Var = [Tkinter.IntVar() for i in range(17)]
        for i in range(1,5):
            self.selected_Var[i].set(1)
        select_buttons = [None]*17
        for i in range(1,17):
            select_buttons[i] = Tkinter.Checkbutton(self.sample_select_box, text=str(i), variable=self.selected_Var[i], command=self.OnSampleSelect)
            select_buttons[i].grid(column=(i-1)%8,row=(i-1)//8)

        self.auto_tc_Var = Tkinter.IntVar()
        self.auto_tc_Var.set(1)
        auto_tc_button = Tkinter.Checkbutton(self, text='Auto interval', variable=self.auto_tc_Var)
        auto_tc_button.grid(column=3,row=5,sticky='W')

        self.time_Label = Tkinter.Label(self, text = "Current time:\nNot running")
        self.time_Label.grid(column=4,row=5,columnspan=2,sticky='W')

        self.start_button_Variable = Tkinter.StringVar()
        self.start_button = ttk.Button(self,textvariable=self.start_button_Variable,command=self.OnButtonClick)
        self.start_button.grid(column=9,row=4,rowspan=2,sticky='EWNS')
        self.start_button_Variable.set('START')

        self.status_bar = thread_safe_tk_widgets.ThreadSafeLabel(self,text="Hello!",anchor="w",fg="white",bg="blue")
        self.status_bar.grid(column=0,row=6,columnspan=10,sticky='EW')

        # logic part
        initial_sample_list = set([i for i in range(1,17) if int(self.selected_Var[i].get())==1])
        self.test = slow_test.slow_test(initial_sample_list, SLOW_LOCKIN, FREQ = float(self.slow_freq_Variable.get()), print_ch = 'Tk',
                                        INTERVAL = float(self.initial_interval_Variable.get()), WAIT_TIME = float(self.wait_time_Variable.get()),
                                        Tk_output = self.output_box, Tk_status = self.status_bar, DMM_addr = DMM)
        self.wait_time_Variable.trace('w',self.OnWaitTimeChange)
        self.initial_interval_Variable.trace('w',self.OnInitialInvervalChange)
        self.auto_tc_Var.trace('w',self.OnAutoTcChange)
        self.is_running = False
        self.UpdateTimeLabel()

    def OnButtonClick(self):
        self.start_button.state(['disabled'])
        if not self.is_running: # not running, then start
            self.start_button_Variable.set('STOP') #1
            self.test.interval = [float(self.initial_interval_Variable.get())] * 17
            self.test.box.wait = float(self.wait_time_Variable.get())
            self.test.sample_list = set([i for i in range(1,17) if int(self.selected_Var[i].get())==1])
            self.test._to_stop = False #4
            self.is_running = True #5
            self.freq_Entry.configure(state='disabled') #6
            self.test.freq = float(self.slow_freq_Variable.get())
            #self.test.initialize()
            self.test_thread = threading.Thread(target = self.test.main_test_loop)
            self.test_thread.start()
        else: # already running, then stop
            self.test._to_stop = True #4
            self.test_thread.join()
            self.is_running = False #5
            self.freq_Entry.configure(state='normal') #6
            self.start_button_Variable.set('START') #1
        self.start_button.state(['!disabled'])

    def OnSampleSelect(self):
        for child in self.sample_select_box.winfo_children():
            child.configure(state='disable')
        new_sample_list = set([i for i in range(1,17) if int(self.selected_Var[i].get())==1])
        if self.is_running:
            add_queue = new_sample_list - self.test.sample_list
            remove_queue = self.test.sample_list - new_sample_list
            for sample in add_queue:
                self.test.add_one_to_measurement(sample)
                #self.output_box.write('add sample %i' % sample)
            for sample in remove_queue:
                self.test.remove_one_from_measurement(sample)
                #self.output_box.write('remove sample %i' % sample)
        for child in self.sample_select_box.winfo_children():
            child.configure(state='normal')

    def OnWaitTimeChange(self,*args):
        self.test.wait_time = float(self.wait_time_Variable.get())

    def OnInitialInvervalChange(self,*args):
        self.test.default_interval = float(self.initial_interval_Variable.get())

    def OnAutoTcChange(self,*args):
        self.test._auto_tc = bool(self.auto_tc_Var.get())

    def EntryValidate(self, P):
        # Disallow anything but float number
        valid = bool(re.match(r'^\d\.?\d*$',P))
        if not valid:
            self.bell()
        return valid

    def UpdateTimeLabel(self):
        if self.test.t0 == None:
            self.time_Label.configure(text='Current time:\nNot Running')
        else:
            self.time_Label.configure(text='Current time:\n%fs'%(time.clock()-self.test.t0))
        self.parent.after(1000,self.UpdateTimeLabel)

class main_window(Tkinter.Tk):
    def __init__(self, parent):
        Tkinter.Tk.__init__(self, parent)
        self.parent = parent
        self.initialize()

    def initialize(self):
        self.fast_sample_Variable = Tkinter.StringVar()
        self.slow_freq_Variable = Tkinter.StringVar()
        self.left_frame = fast_frame(self,self.fast_sample_Variable)
        self.left_frame.grid(row=0, column=0)
        self.right_frame = slow_frame(self,self.slow_freq_Variable)
        self.right_frame.grid(row=0, column=1)
        self.title('Multiplexer Measurements')
        self.resizable(False,False)
        self.lift()
        self.focus_force()

        self.protocol("WM_DELETE_WINDOW", self.OnClosingWindow)
        self.fast_sample_Variable.trace('w',self.OnFastSampleChange)
        self.slow_freq_Variable.trace('w',self.OnSlowFreqChange)

    def OnFastSampleChange(self,*args):
        self.right_frame.test.idx_to_ignore = int(self.fast_sample_Variable.get())

    def OnSlowFreqChange(self,*args):
        self.left_frame.test.SLOW_FREQ = float(self.slow_freq_Variable.get())

    def OnClosingWindow(self):
        if self.left_frame.is_running or self.right_frame.is_running:
            tkMessageBox.showerror("Error","Measurement still running")
        else:
            if tkMessageBox.askokcancel("Quit", "Do you want to quit?"):
                root.destroy()

if __name__ == "__main__":
    logging.basicConfig(filename = 'window_errors.log', level=logging.WARNING)
    root = main_window(None)
    root.mainloop()