import tkinter as tk
from tkinter import ttk
from tkinter import filedialog
import os
import glob
import cv2
import math
import csv
import re
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from tensorflow.keras.applications.imagenet_utils import preprocess_input, decode_predictions
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 import ResNet50
from keras.utils import np_utils
from skimage.transform import resize
from sklearn.model_selection import train_test_split
from tensorflow.keras.layers import Dense, InputLayer, Dropout
from tensorflow.keras.models import Sequential

def load_train_tab(tab_train, root):
    def browse_directory():
        global selected_directory
        filepath = filedialog.askdirectory(title="Select a directory")
        if filepath:
            print("Selected file:", filepath)
            selected_directory = filepath # Save the directory name to the global variable
            directory_entry.delete(0, tk.END)
            directory_entry.insert(0, filepath)

            # Update the video files list
            video_files.config(state="normal")  # Set the state to normal before inserting text
            video_files.delete(1.0, tk.END)
            for file in os.listdir(filepath):
                if file.endswith(".mp4") or file.endswith(".avi") or file.endswith(".mkv"):
                    video_files.insert(tk.END, file + "\n")
            video_files.config(state="disabled")  # Set the state back to disabled after inserting text

    def frame_split():
        # Create a progress bar popup window
        progress_window = tk.Toplevel(root)
        progress_window.title("Splitting Videos")
        progress_label = ttk.Label(progress_window, text="Splitting video into frames...")
        progress_label.pack(padx=10, pady=(10, 0))
        progress = ttk.Progressbar(progress_window, mode="indeterminate", length=300)
        progress.pack(padx=10, pady=(5, 10))
        progress.start(10)
        progress_window.update()

        #check if train/frames path exists, if not, create it
        if os.path.exists(selected_directory + "/train/frames/") == False:
            print("/train/frames folder does not exist. Creating...")
            os.makedirs(selected_directory + "/train/frames/")
        else:
            print("train/frames folder already exists")

        #capture video files in chosen directory
        count = 0
        cap = [cv2.VideoCapture(videoFile) for videoFile in glob.glob(os.path.join(selected_directory, "*.mp4"))] # capturing the video from the given path

        #split the frames from each video then output to train/frames folder
        for i in cap:
            print(str(i))
            frameRate = i.get(5)
            while (i.isOpened()):
                frameID = i.get(1)
                ret, frame = i.read()
                if (ret != True):
                    break
                if (frameID % math.floor(frameRate) == 0):
                    filename = selected_directory + "/train/frames/frame%d.jpg" % (count); count +=1
                    cv2.imwrite(filename, frame)
            i.release()

        #create the excel file from split frames
        print("Creating excel file for classification...")
        header = ['Image_ID', 'Class']
        data = []
        for i in os.listdir(selected_directory + "/train/frames"):
            data.append(str(i))
        data.sort(key=lambda f: int(re.sub('\D', '', f)))
        data2 = []
        for i in data:
            data2.append([i])
        with open(selected_directory + '/train/frames.csv', 'w', encoding='UTF8', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(header)
            writer.writerows(data2)
        print("Done! Please label frames accordingly in 'frames.csv' ")

        # Close the progress bar window
        progress.stop()
        progress_window.destroy()

        # Show a new popup window that says "frame split complete"
        complete_window = tk.Toplevel(root)
        complete_window.title("Complete")
        complete_label = ttk.Label(complete_window, text="Frame splitting complete. \nYour training frames are located in /train/frames/ in your selected directory. \nPlease update your excel file located in the /train/ folder with the necessary labels")
        complete_label.pack(padx=10, pady=(10, 0))
        ok_button = ttk.Button(complete_window, text="OK", command=complete_window.destroy)
        ok_button.pack(padx=10, pady=(5, 10))

        # Update the main window
        root.update()

    def begin_training():
        # Create a progress bar popup window
        progress_window = tk.Toplevel(root)
        progress_window.title("Splitting Videos")
        progress_label = ttk.Label(progress_window, text="Training model...")
        progress_label.pack(padx=10, pady=(10, 0))
        progress = ttk.Progressbar(progress_window, mode="indeterminate", length=300)
        progress.pack(padx=10, pady=(5, 10))
        progress.start(10)
        progress_window.update()

        #load training excel file
        data = pd.read_csv(selected_directory + '/train/frames.csv')

        #count number of areas created in excel file:
        cnt = 0
        visited = []
        for i in range (0, len(data['Class'])):
            if data['Class'][i] not in visited:
                visited.append(data['Class'][i])

                cnt+=1
        X = [ ]     # creating an empty array
        for img_name in data.Image_ID:
            img = plt.imread(selected_directory + '/train/frames/' + img_name)
            X.append(img)  # storing each image in array X
        X = np.array(X)
        y = data.Class
        dummy_y = np_utils.to_categorical(y)
        image = []
        for i in range(0,X.shape[0]):
            a = resize(X[i], preserve_range=True, output_shape=(224,224)).astype(int)      # reshaping to 224*224*3
            image.append(a)
        X = np.array(image)
        X = preprocess_input(X, mode='caffe')
        X_train, X_valid, y_train, y_valid = train_test_split(X, dummy_y, test_size=0.3, random_state=42)
        base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
        X_train = base_model.predict(X_train)
        X_valid = base_model.predict(X_valid)
        X_train.shape, X_valid.shape
        X_train = X_train.reshape(X_train.shape[0], 7*7*2048)      # converting to 1-D
        X_valid = X_valid.reshape(X_valid.shape[0], 7*7*2048)
        train = X_train/X_train.max()      # centering the data
        X_valid = X_valid/X_train.max()
        model = Sequential()
        model.add(InputLayer((7*7*2048,)))    # input layer
        model.add(Dense(units=2048, activation='sigmoid'))   # hidden layer
        model.add(Dropout(0.5))      # adding dropout
        model.add(Dense(units=1024, activation='sigmoid'))    # hidden layer
        model.add(Dropout(0.5))      # adding dropout
        model.add(Dense(units=512, activation='sigmoid'))    # hidden layer
        model.add(Dropout(0.5))      # adding dropout
        model.add(Dense(units=256, activation='sigmoid'))    # hidden layer
        model.add(Dropout(0.5))      # adding dropout
        model.add(Dense(cnt, activation='softmax'))            # output layer
        model.summary()
        model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
        model.fit(train, y_train, epochs=15, validation_data=(X_valid, y_valid))
        model.save(selected_directory + '/train/model')
        # Close the progress bar window
        progress.stop()
        progress_window.destroy()

        # Show a new popup window that says "model training complete"
        complete_window = tk.Toplevel(root)
        complete_window.title("Complete")
        complete_label = ttk.Label(complete_window, text="Model training complete. Model has been saved to /train/model/.\nYou may begin classification of new videos in the classify tab")
        complete_label.pack(padx=10, pady=(10, 0))
        ok_button = ttk.Button(complete_window, text="OK", command=complete_window.destroy)
        ok_button.pack(padx=10, pady=(5, 10))

        # Update the main window
        root.update()

    # Add description text above the browse button
    description_label = ttk.Label(tab_train, text="Select user video folder:")
    description_label.pack(padx=10, pady=(10, 0))

    # Create a frame to hold the directory entry and browse button
    entry_browse_frame = ttk.Frame(tab_train)
    entry_browse_frame.pack(padx=10, pady=(5, 10))

    # Add an empty text box for manual directory input
    directory_entry = ttk.Entry(entry_browse_frame, width=50)
    directory_entry.pack(side="left")

    # Add a browse button to the "Train Model" tab
    browse_button = ttk.Button(entry_browse_frame, text="Browse", command=browse_directory)
    browse_button.pack(side="left", padx=(10, 0))

    # Create a text box to show a list of video files in the chosen directory
    video_files = tk.Text(tab_train, wrap="none", width=50, height=10, state="normal")
    video_files.pack(padx=10, pady=(5, 10))

    # Add split text description above the split button
    split_description_label = ttk.Label(tab_train, text="Split videos into frames and output to a train folder:")
    split_description_label.pack(padx=10, pady=(10, 0))

    # Add a split button to the "Train Model" tab
    split_button = ttk.Button(tab_train, text="Split", command=frame_split)
    split_button.pack(padx=10, pady=(5, 10))

    #Add train text description above train button
    train_description_label = ttk.Label(tab_train, text="Begin training model (please make sure your excel file is properly filled out)")
    train_description_label.pack(padx=10, pady=(10, 0))

    # Add a train button to the "Train Model" tab
    train_button = ttk.Button(tab_train, text="Train", command=begin_training)
    train_button.pack(padx=10, pady=(5, 10))


import cantera as ct
import numpy as np
from typing import List, Tuple
from scipy import integrate
from copy import copy

"""
Present a simple implementation of IDT reactors and the
cantera implementation of a LFS reactor.

Each model can be called as:

IDT, all_conditions = idt_reactor.solve(gas, flag='T', temp_rise=400)
IDT, all_conditions = idt_reactor.solve(gas, path_to_save=dir_to_save, phi_for_file_name=phi_value)
WARNINGS: idt_reactor uses only the Temperarure rise. The species peak is still under developement.

LFS, all_conditions = lfs_reactor.solve(gas)
LFS, all_conditions = lfs_reactor.solve(gas, path_to_save=dir_to_save, phi_for_file_name=phi_value)

The saved file will be named as:
<model_type>_<TEMP in K>_<PRESSURE in atm>_<PHI, default is 1.0>.csv

The first line will contain the propertie value (IDT or LFS).
The second line will be a header containing specifics and the conditions.
The rest of the file will present the values in a csv format.
"""

# --------------------------------------------------------------------------------------------------
# utils for save
def save_solution_to_file(file_path: str,
                 model_type: str,
                 model_value: float,
                 header: str,
                 states: np.ndarray,) -> None:
    """
    Save the conditions to a csv file.

    The fist line presents the model propertie, folowed by the header and the data.

    MODEL_<model specific>=<model value>
    MODEL_IDT_T_400 -> IDT model, with T flag as 400 K
    MODEL_LFS -> LFS model
    """
    data_to_save = f"MODEL_{model_type}={model_value}\n{header}\n"
    data_to_save += "\n".join([",".join([str(states[row, col]) for col in range(states.shape[1])]) for row in range(states.shape[0])])

    with open(file_path, 'w') as f:
        f.write(data_to_save)

def create_solution_file_name(model_type: str, temp: float, press: float, phi: float) -> str:
    """
    Creates the file name considering the initial state.

    <model_type>_<temp>_<press>_<phi>.csv

    temp -> Kelvin
    press -> Pa

    if no phi is provided, use 1.0 as default.
    """
    f_name = f"{model_type}_{temp:.1f}_{press/ct.one_atm:.2f}_"
    if phi:
        f_name += f"{phi:.1f}"
    else:
        f_name += "1.0"
    f_name += ".csv"
    return f_name

# --------------------------------------------------------------------------------------------------
# IDT implementation
class ConstantMassPressureODE:
    """Implement the 0D, constant mass, constant pressure reactor"""
    def __init__(self, gas: ct.Solution) -> None:
        self.gas: ct.Solution = gas
        self._pressure: float = gas.P

    def __call__(self, t: float, y: np.ndarray) -> np.ndarray:
        """return the set of EDO. See Turns to understand."""
        # Set the gas conditions
        if y[0] <= 0:
            raise ValueError(f"Negative value found for temperature.")
        self.gas.set_unnormalized_mass_fractions(y[1:])
        self.gas.TP = y[0], self._pressure

        # calculates all the values
        rho = self.gas.density
        wdot = self.gas.net_production_rates
        dTdt = - (np.dot(self.gas.partial_molar_enthalpies, wdot) / (rho * self.gas.cp))
        dYdt = wdot * self.gas.molecular_weights / rho
        return np.hstack((dTdt, dYdt))

class IDTReactor:
    """
    Class implementation of an 0D reactor to obtain the IDT value.
    """
    def __init__(self, rtol: float = 10**-6, atol: float = 10**-9, n_iter: int = 10000) -> None:
        self.rtol = rtol
        self.atol = atol
        self.n_iter = n_iter

    def _get_solver(self, gas: ct.Solution, inital_condition: np.ndarray) -> integrate.ode:
        """
        Set the solver to run the IDT cases.
        """
        ODE = ConstantMassPressureODE(gas)
        solver = integrate.ode(ODE)
        solver.set_integrator('lsoda', method='bdf', rtol=self.rtol, atol=self.atol, with_jacobian=False, nsteps=self.n_iter)
        solver.set_initial_value(inital_condition, 0.0)
        return solver

    def _get_header(self, gas: ct.Solution) -> str:
        return "time(s),T,P," + ",".join(["X_" + spc.name for spc in gas.species()])

    def get_idt(self, gas: ct.Solution,
                    max_time: float = 5.0,
                    dt: float = 10**-7,
                    flag: str = 'T',
                    temp_rise: float = 400.0) -> float:
        """
        Find the idt time.

        This returns a float. If an error is raised by solver problems, it returns a -1.0,
        if no IDT is found in the time window, returns -2.0.
        """
        # TODO: Add the species flag. For a peak, run up to the max time.

        # Prepare the initial conditions
        initial_condition = np.array([gas.T, *gas.Y])
        temperature_flag = temp_rise + initial_condition[0]

        # Set the solver
        solver = self._get_solver(gas, initial_condition)

        # solve
        try:
            while solver.successful() and solver.t <= max_time:
                if solver.y[0] >= temperature_flag:
                    return solver.t
                solver.integrate(solver.t + dt)

        # catch any temperature problem
        except:
            return -1.0

        # if we do not find a IDT in the max_time
        return -2.0

    def get_norm_states(self, gas: ct.Solution, idt: float,
                            norm_dt: float = 0.01,
                            max_norm_time: float = 2.0) -> Tuple[np.ndarray, str]:
        """
        Solve the idt reactor at every norm dt and return the real time conditions.

        Returns a np.ndarray containing the values and a str containig the header:
        time(s), T, P, Y_<spc names>
        """
        initial_condition = np.array([gas.T, *gas.Y])
        const_pressure = copy(gas.P)
        n_points = int(max_norm_time / norm_dt + 1)
        out_solution = np.zeros([n_points, 3 + gas.n_species])
        out_solution[0, 1], out_solution[0, 2]  = gas.T, gas.P
        out_solution[0, 3:] = gas.X

        # Set the solver
        solver = self._get_solver(gas, initial_condition)

        # set control parameters
        dt = norm_dt * idt
        max_time = max_norm_time * idt
        cur_point = 0

        try:
            while solver.successful():
                cur_point += 1
                solver.integrate(solver.t + dt)

                # for the output to be in molar fraction
                gas.TPY = solver.y[0], const_pressure, solver.y[1:]
                out_solution[cur_point, 0] = solver.t
                out_solution[cur_point, 1] = gas.T
                out_solution[cur_point, 2] = gas.P
                out_solution[cur_point, 3:] = gas.X

                if cur_point == n_points - 1:
                    break

            return out_solution, self._get_header(gas)
        except:
            raise Exception("Failed to solve the ODE. Try a different set of tolerances.")

    def solve(self, gas: ct.Solution,
                    path_to_save: str = "",
                    phi_for_file_name: float = None,
                    norm_dt: float = 0.01,
                    max_norm_time: float = 2.0,
                    max_time_for_idt: float = 5.0,
                    dt_for_idt: float = 10**-7,
                    idt_flag: str = 'T',
                    idt_temp_rise: float = 400.0) -> Tuple[float, np.ndarray, str]:
        """
        Solve the reactor, returning a IDT value, a np.ndarray with all conditions and the corresponding header.
        If a directory is passed as input, save the conditions to the file.

        The condition in the gas solution passed to this method is consdered the initial condition.
        """
        init_TPY = copy(gas.TPY)
        idt = self.get_idt(gas, max_time=max_time_for_idt, dt=dt_for_idt, flag=idt_flag, temp_rise=idt_temp_rise)

        if idt <= 0.0:
            if idt == -2.0:
                raise Exception(f"It was not possble to obtain IDT. No IDT found in the time window.")
            raise Exception(f"It was not possble to obtain IDT. Solver problems found.")

        gas.TPY = init_TPY
        states, header = self.get_norm_states(gas, idt, norm_dt=norm_dt, max_norm_time=max_norm_time)

        # check for save flag:
        if path_to_save != "":
            f_name = create_solution_file_name("IDT", init_TPY[0], init_TPY[1], phi_for_file_name)
            save_solution_to_file(path_to_save + f_name, f"IDT_{idt_flag}_{idt_temp_rise:.2f}", idt, header, states)

        return idt, states, header

# --------------------------------------------------------------------------------------------------
# LFS implementation
class LFSReactor:
    """
    Class implementation of an 1D reactor to obtain the LFS value.
    """
    def __init__(self, width: float = 0.014,
                        ratio: float=3,
                        slope: float=0.1,
                        curve: float=0.1,
                        max_time_step_count: int = 5000,
                        loglevel: int = 0) -> None:
        self.width = width
        self.ratio = ratio
        self.slope = slope
        self.curve = curve
        self.max_time_step_count = max_time_step_count
        self.loglevel = loglevel

    def _get_header(self, gas: ct.Solution) -> str:
        return "grid(s),T,P," + ",".join(["X_" + spc.name for spc in gas.species()])

    def _get_states(self, flame: ct.FreeFlame) -> np.ndarray:
        """
        Return the states of the current flame.

        grid(m), T, P, X_<species>
        """
        out_data = np.zeros([len(flame.T),len(flame.X) + 3])
        out_data[:,0] = flame.grid
        out_data[:,1] = flame.T
        out_data[:,2] = flame.P
        for i in range(len(flame.X)):
            out_data[:,3 + i] = flame.X[i]
        return out_data

    def solve(self, gas: ct.Solution,
                    path_to_save: str = "",
                    phi_for_file_name: float = None) -> Tuple[float, np.ndarray, str]:
        """
        Solve the reactor, returning a IDT value, a np.ndarray with all conditions and the corresponding header.
        If a directory is passed as input, save the conditions to the file.

        The condition in the gas solution passed to this method is consdered the initial condition.
        """
        init_TPY = copy(gas.TPY)

        # Create the flame object
        flame = ct.FreeFlame(gas, width=self.width)
        # flame.transport_model = 'Mix'

        # Define tolerances for the solver
        flame.set_refine_criteria(ratio=self.ratio, slope=self.slope, curve=self.curve)
        flame.max_time_step_count = self.max_time_step_count

        # Define logging level
        flame.solve(loglevel=self.loglevel, auto=True)
        Su0 = flame.velocity[0]
        states = self._get_states(flame)
        header = self._get_header(gas)

        # check for save flag:
        if path_to_save != "":
            f_name = create_solution_file_name("LFS", init_TPY[0], init_TPY[1], phi_for_file_name)
            save_solution_to_file(path_to_save + f_name, f"LFS_m_s", Su0, header, states)

        return Su0, states, header


from scrapy.selector import HtmlXPathSelector
from scrapy.contrib.spiders import CrawlSpider  # Rule
from scrapy.http.request import Request
import html2text
import time
import re
import dateutil.parser
import datetime
import urlparse
from buzz_crawler.items import BuzzCrawlerItem
from markdown import markdown

class WozSpider(CrawlSpider):
    name = 'woz'
    allowed_domains = ['www.woz.ch']
    start_urls = ['http://www.woz.ch/']

    def handle_blog(self, response):
        hxs = HtmlXPathSelector(response)
        item = BuzzCrawlerItem()

        item['url'] = response.url

        item['date'] = datetime.datetime.now()
        item['title'] = hxs.xpath(".//*[@id='container']/div/div/article/header/h1/text()").extract()[0].strip()
        item['blurb'] = hxs.xpath(".//*[@id='container']/div/div/article/header/h2/text()").extract()[0].strip()

        unprocessed_content = hxs.xpath("//div[@class='article-content']").extract()[0].strip()
        h = html2text.HTML2Text()
        h.ignore_links = True
        h.ignore_images = True
        processed_content = h.handle(unprocessed_content)

        item['content'] = markdown(processed_content)
        item['source'] = 'woz.ch'
        yield item

    def parse(self, response):
        hxs = HtmlXPathSelector(response)
        posts = hxs.xpath(".//*[@id='container']/div/div/article")

        for post in posts:
            post_link = post.xpath("a/@href").extract()[0]
            post_absolute_url = urlparse.urljoin(response.url, post_link.strip())
            yield Request(post_absolute_url, self.handle_blog)


import xml.etree.ElementTree as ET
import json

path_train = "D:/code/prompt-ABSA/dataset/original data/ABSA16_Laptops_Train_SB1_v2.xml"
path_test = 'D:/code/prompt-ABSA/dataset/original data/EN_LAPT_SB1_TEST_.xml'
terr = 'laptops16'

def get_path(territory, data_type):
    return f'./dataset/data/{territory}/{data_type}.json'

def judge(p):
    for i in range(len(p)):
        for j in range(i + 1, len(p)):
            if p[j]['category'] == p[i]['category'] and p[j]['polarity'] != p[i]['polarity']:
                return False
    return True

def extract_data(path):
    tree = ET.parse(path)
    root = tree.getroot()
    data = []

    # 从xml文件中提取数据
    for review in root:
        for sentences in review:
            for sentence in sentences:
                piece = []
                for t in sentence.iter('text'):
                    piece.append(t.text)
                for o in sentence.iter('Opinion'):
                    d = {'category': o.attrib['category'], 'polarity': o.attrib['polarity']}
                    piece.append(d)
                # 所有沒有category分類的句子以及所有一個category卻多個情感的句子
                if len(piece) > 1 and judge(piece[1:]):
                    data.append(piece)

    n_category = 0
    category = []

    # 进行数据统计
    for e in data:
        for i in range(1, len(e)):
            c, s = e[i].values()
            if c not in category:
                n_category += 1
                category.append(c)

    all_data = [n_category, category, data]
    return all_data

train_data = extract_data(path_train)
test_data = extract_data(path_test)
# 将train中没有而test中有的category從test中刪除

over_list = [elem for elem in test_data[1] if elem not in train_data[1]]
move_list = [elem for cate in over_list for elem in test_data[2] for e in elem[1:] if e['category'] == cate]
test_data[2] = [elem for elem in test_data[2] if elem not in move_list]
test_data[1] = [elem for elem in test_data[1] if elem not in over_list]
test_data[0] = len(test_data[1])
print(over_list)

with open(get_path(terr, 'train'), 'w', encoding='utf-8') as f:
    json.dump(train_data, f)
with open(get_path(terr, 'test'), 'w', encoding='utf-8') as f:
    json.dump(test_data, f)


"""Basic state machine implementation."""
# pylint: disable=unnecessary-pass, too-many-instance-attributes

from typing import Iterable, Union
from rclpy import logging
from rclpy.node import Node
from rclpy.time import Time, Duration

LOGGER = logging.get_logger("behavior")

class Resource:
    """The resource class is used to track access to shared resources.

    When a state starts, it attempts to acquire all the resources it needs.
    Resources are automatically freed when the state stops. By default,
    multiple states can access the same resource, but a UniqueResource will
    throw an error if a state tries to acquire a resource that is already owned
    by another state.

    The registry is a class variable that is used to keep track of all the
    resources and their owners.
    """
    _registry = {}

    def __init__(self, name: str):
        self._name = name
        if name in Resource._registry:
            LOGGER.error(f"Resource {name} already exists. "
                "Different resources must have unique names.")
        Resource._registry[name] = set()

    def __del__(self):
        Resource._registry.pop(self._name)

    def get_name(self) -> str:
        """Get the name of the resource"""
        return self._name

    def get_owner(self):
        """Get the owner of the resource"""
        return Resource._registry[self._name]

    def is_free(self) -> bool:
        """Check if the resource is free"""
        return len(Resource._registry[self._name]) == 0

    def assign(self, owner: 'State') -> 'Resource':
        """Assign the resource to an owner."""
        if owner in Resource._registry[self._name]:
            LOGGER.error(f"Resource {self._name} already owned by {owner}.")
        Resource._registry[self._name].add(owner)
        return self

    def transfer(self, old_owner: 'State', new_owner: 'State') -> 'Resource':
        """Transfer ownership of the resource"""
        if old_owner not in Resource._registry[self._name]:
            LOGGER.error(f"Resource {self._name} not owned by {old_owner}.")
        Resource._registry[self._name].remove(old_owner)
        Resource._registry[self._name].add(new_owner)
        return self

    def free(self, owner: 'State'):
        """Free the resource"""
        Resource._registry[self._name].remove(owner)

    @staticmethod
    def validate(*resources: str):
        """Check if a set of resource names is valid"""
        for resource in resources:
            if resource not in Resource._registry:
                LOGGER.error(f"Resource {resource} does not exist. Invalid name.")

class UniqueResource(Resource):
    """A resource that can only be owned by one state at a time"""
    def assign(self, owner: 'State') -> Resource:
        """Assign the resource to an owner."""
        if not self.is_free():
            LOGGER.error(f"Resource {self._name} already owned by "
                         f"{Resource._registry[self._name]}.")
        Resource._registry[self._name].add(owner)
        return self

class State:
    """Superclass for states"""
    def __init__(self, name,
                 node:Node = None,
                 resources:Iterable[str] = None,
                 parent:'StateMachine' = None):
        self._base_name = name
        self._name = name
        self._node = None
        self._running = False
        self._start_time = None
        self._required_resources = set() if resources is None else set(resources)
        self._parent = None
        self._resources = {}
        if node is not None:
            self._node = node
        if parent is not None:
            self.set_parent(parent)

    def set_node(self, node:Node):
        """Set the node for the state"""
        self._node = node

    def set_parent(self, parent:'StateMachine'):
        """Set the parent state machine"""
        if self._parent is not None:
            LOGGER.error(f"State {self._name} already has parent.")
        if self._running:
            LOGGER.error(f"State {self._name} already running. "
                         "Cannot change parent.")
        self._parent = parent
        self._name = f"{parent.get_name()}/{self._base_name}"

    def get_base_name(self):
        """Get the base name of the state"""
        return self._base_name

    def get_name(self):
        """Get the name of the state"""
        return self._name

    def get_required_resources(self) -> Iterable[str]:
        """Get the set of required resources"""
        return self._required_resources

    def add_required_resource(self, resource: Union[str, Iterable[str]]):
        """Add a required resource"""
        if self._running:
            LOGGER.error(f"State {self._name} already running. "
                         "Cannot add required resources.")
        if isinstance(resource, str):
            self._required_resources.add(resource)
        else:
            for res in resource:
                self._required_resources.add(res)
        if self._parent is not None:
            self._parent.add_required_resource(resource)

    def get_start_time(self) -> Time:
        """Get the time the state started"""
        return self._start_time

    def get_current_time(self) -> Time:
        """Get the current time"""
        return self._node.get_clock().now()

    def get_elapsed_time(self) -> Duration:
        """Get the time the state has been running"""
        return self._node.get_clock().now() - self._start_time

    def get_node(self) -> Node:
        """Get the node for the state"""
        return self._node

    def start(self):
        """Start the state"""
        if self._node is None:
            LOGGER.error(f"State {self._name} must be run within a ros2 node. "
                         "Use set_node() to set the node.")
        if self._running:
            LOGGER.warning(f"State {self._name} already running")
        for resource in self._required_resources:
            self._acquire_resource(resource)
        self._running = True
        self._start_time = self._node.get_clock().now()

    def stop(self):
        """Stop the state"""
        if not self._running:
            LOGGER.warning(f"State {self._name} not running")
            return
        self._release_all_resources()
        self._running = False

    def step(self):
        """Do one step of the state"""
        if not self._running:
            LOGGER.error(f"State {self._name} has not been started. Cannot step.")

    def get_resource(self, resource: str) -> Resource:
        """Get a resource reference"""
        return self._resources[resource]

    def _acquire_resource(self, resource_name: str):
        """Acquire a resource"""
        if resource_name in self._resources:
            LOGGER.warning(f"Resource {resource_name} already acquired")
        resource = self._parent.get_resource(resource_name)
        if isinstance(resource, UniqueResource):
            self._resources[resource_name] = resource.transfer(self._parent, self)
        else:
            self._resources[resource_name] = resource.assign(self)

    def _release_resource(self, resource_name: str):
        """Release a resource back to parent"""
        if resource_name not in self._resources:
            LOGGER.error(f"Resource {resource_name} not in resource list for "
                         f"{self.get_name()}.")
        resource = self._resources[resource_name]
        if isinstance(resource, UniqueResource):
            resource.transfer(self, self._parent)
        else:
            resource.free(self)
        self._resources[resource].transfer(self, self._parent)
        self._resources.pop(resource)

    def _release_all_resources(self):
        """Release all resources"""
        for resource in self._resources.values():
            if isinstance(resource, UniqueResource):
                resource.transfer(self, self._parent)
            else:
                resource.free(self)
        self._resources.clear()

    def validate(self):
        """Validate the state"""
        if self._parent is None:
            # Only need to validate resource names at global level
            Resource.validate(*self._required_resources)

class Event:
    """An event that signals a transtition between states"""

    def __init__(self, name):
        self._name = name
        self._status = False
        self._enabled = False
        self._state = None
        self._required_resources = set()
        self.initialized = True # Some events might need additional info before
                                # being initialized such as subscriber events

    def get_name(self) -> str:
        """Get the name of the event"""
        return self._name

    def reset(self):
        """Reset the event"""
        self._status = False

    def get_status(self) -> bool:
        """Get the status of the event without updating"""
        return self._status

    def update(self) -> bool:
        """Update the event and get status"""
        return self._status

    def add_required_resource(self, resource: Union[str, Iterable[str]]):
        """Add a required resource"""
        if isinstance(resource, str):
            self._required_resources.add(resource)
        else:
            for res in resource:
                self._required_resources.add(res)

    def get_required_resources(self) -> Iterable[str]:
        """Get the set of required resources"""
        return self._required_resources

    def enable(self, state:State):
        """Called when corresponding state is started"""
        if self._state is not None:
            LOGGER.error(f"Event {self._name} already enabled for state "
                         f"{self._state.get_name()}")
        self._state = state
        self._enabled = True

    def is_enabled(self) -> bool:
        """Check if event is enabled"""
        return self._enabled

    def disable(self):
        """Called when corresponding state is stopped"""
        self.reset()
        self._state = None
        self._enabled = False

    def __invert__(self) -> 'Event':
        if isinstance(self, NotEvent):
            return self._e1
        return NotEvent(self)

    def __and__(self, other:'Event') -> 'Event':
        return AndEvent(self, other)

    def __or__(self, other:'Event') -> 'Event':
        return OrEvent(self, other)

class CompositeEvent(Event):
    """An event that is a combination of other events."""
    def __init__(self, name):
        super().__init__(name)
        self._events = set()

    def get_base_events(self) -> Iterable[Event]:
        """Returns the base (non-composite) events"""
        return self._events

    def add_event(self, event:'Event'):
        """Add an event to the set"""
        if isinstance(event, self.__class__):
            self._events |= event.get_base_events()
        self._events.add(event)
        self._required_resources |= event.get_required_resources()

    def reset(self):
        for event in self._events:
            event.reset()
        self._status = False

    def update(self) -> bool:
        for event in self._events:
            event.update()
        return self._status

    def enable(self, state:State):
        super().enable(state)
        for event in self._events:
            if not event.is_enabled():
                event.enable(state)

class NotEvent(CompositeEvent):
    """An event that is true if its base event is false"""
    def __init__(self, event, name=None):
        super().__init__(f"not_{event.get_name()}" if name is None else name)
        self.add_event(event)
        self._e1 = event

    def update(self) -> bool:
        if not self._e1.initialized:
            return False
        self._status = not self._e1.update()
        return self._status

class AndEvent(CompositeEvent):
    """An event that is true if both of two events is true"""
    def __init__(self, event1, event2, name=None):
        super().__init__(f"{event1.get_name()}_and_{event2.get_name()}"
                         if name is None else name)
        self.add_event(event1)
        self.add_event(event2)
        self._e1 = event1
        self._e2 = event2

    def update(self) -> bool:
        super().update()
        self._status = self._e1.get_status() and self._e2.get_status()
        return self._status

class OrEvent(CompositeEvent):
    """An event that is true if at least one of two events is true"""
    def __init__(self, event1, event2, name=None):
        super().__init__(f"{event1.get_name()}_or_{event2.get_name()}"
                         if name is None else name)
        self.add_event(event1)
        self.add_event(event2)
        self._e1 = event1
        self._e2 = event2

    def update(self) -> bool:
        super().update()
        self._status = self._e1.get_status() or self._e2.get_status()
        return self._status

class StateMachine(State):
    """A basic state machine"""
    def __init__(self, name, node:Node = None):
        super().__init__(name, node=node)
        self._states = {} # type: Dict[str,State]
        self._events = {} # type: Dict[str,Event]
        self._transitions = {} # type: Dict[str,Dict[str,str]] # from -> event -> to
        self._current_state = None
        self._start_state = None

    def set_node(self, node:Node):
        super().set_node(node)
        for state in self._states.values():
            state.set_node(node)

    def add_state(self, state:State):
        """Add a state to the state machine"""
        if state.get_base_name() in self._states:
            LOGGER.warning(f"State {state.get_base_name()} already in state "
                            "machine. Skipping.")
            return
        state.set_node(self._node)
        state.set_parent(self)
        self._states[state.get_base_name()] = state
        self._required_resources |= state.get_required_resources()
        # By default, the start state is the first state added
        if self._start_state is None:
            self._start_state = state.get_base_name()

    def add_transition(self,
                       from_state:Union[str,State],
                       event:Event,
                       to_state:Union[str,State]):
        """Add a transition to the state machine"""
        if isinstance(from_state, State):
            from_state = from_state.get_base_name()
        if isinstance(to_state, State):
            to_state = to_state.get_base_name()
        if from_state not in self._states:
            LOGGER.error(f"State {from_state} not in state machine. Invalid transition.")
        if to_state not in self._states:
            LOGGER.error(f"State {to_state} not in state machine. Invalid transition.")
        if from_state not in self._events:
            self._events[from_state] = {}
        elif event in self._events[from_state]:
            LOGGER.warning(f"Transition from {from_state} on event {event} already exists. "
                           "Overwriting.")
        self._events[event.get_name()] = event
        if from_state not in self._transitions:
            self._transitions[from_state] = {}
        self._transitions[from_state][event.get_name()] = to_state
        self._states[from_state].add_required_resource(event.get_required_resources())

    def set_start(self, state:Union[str,State]):
        """Set the start state"""
        if isinstance(state, State):
            state = state.get_base_name()
        if state not in self._states:
            LOGGER.error(f"State {state} not in state machine. Invalid start state.")
        self._start_state = state

    def _transition(self, state:Union[str,State]):
        """Transition to a new state"""
        if isinstance(state, State):
            state = state.get_base_name()
        if state not in self._states:
            LOGGER.error(f"State {state} not in state machine. Invalid transition.")
        if self._current_state is not None:
            for event in self.get_all_base_events():
                event.disable()
            self._current_state.stop()
        self._current_state = self._states[state]
        self._current_state.start()
        if self._current_state.get_base_name() in self._transitions:
            for event_name in self._transitions[self._current_state.get_base_name()]:
                event = self._events[event_name]
                event.enable(self._current_state)
        LOGGER.debug(f"Transitioned to state {state}")
        print (f"======Transitioned to state {state}======")

    def get_all_base_events(self, state=None) -> Iterable[Event]:
        """Recursively gets a set of all base events for a state"""
        if state is None:
            state = self._current_state
        events = set()
        if state.get_base_name() not in self._transitions:
            return events
        for event_name in self._transitions[state.get_base_name()]:
            event = self._events[event_name]
            if isinstance(event, CompositeEvent):
                events |= self._events[event_name].get_base_events()
            events.add(event)
        return events

    def start(self):
        """Start the state machine"""
        super().start()
        if self._start_state is None:
            return
        self._transition(self._start_state)
        LOGGER.debug(f"Started state machine {self._name}")

    def stop(self):
        """Stop the state machine"""
        if self._current_state is not None:
            for event in self.get_all_base_events():
                event.disable()
            self._current_state.stop()
        super().stop()
        LOGGER.debug(f"Stopped state machine {self._name}")

    def step(self):
        """Do one step of the state machine"""
        if self._current_state is None:
            return
        if self._current_state.get_base_name() in self._transitions:
            for event_name, to_state in self._transitions[
                    self._current_state.get_base_name()].items():
                if self._events[event_name].update():
                    self._node.get_logger().info(
                        f"{event_name}: {self._current_state.get_base_name()} -> {to_state}")
                    self._transition(to_state)
                    break
        self._current_state.step()

    def _check_reachability(self, start_state:str, reachable:set):
        """Recursively check reachability of states"""
        if start_state in reachable:
            return
        reachable.add(start_state)
        if start_state not in self._transitions:
            return
        for to_state in self._transitions[start_state].values():
            self._check_reachability(to_state, reachable)

    def validate(self):
        """Validates State machine consturction to check for errors at initialization"""
        super().validate()
        if self._start_state is None:
            return
        # Recursively validate child state machines
        for state in self._states.values():
            state.validate()
        # Check reachability
        reachable = set()
        self._check_reachability(self._start_state, reachable)
        unreachable = set(self._states.keys()) - reachable
        if len(unreachable) > 0:
            LOGGER.warning(f"State machine {self._name} has unreachable states: {unreachable}")
        # Check resource sharing
        for state in self._states.values():
            missing_resources = state.get_required_resources() - self._required_resources
            if len(missing_resources) > 0:
                LOGGER.error(f"State {state.get_name()} requires resources {missing_resources} "
                             f"that are not provided by state machine {self._name}")
        for state, events in self._transitions.items():
            for event in events:
                missing_resources = (self._events[event].get_required_resources() -
                                     self._states[state].get_required_resources())
                if len(missing_resources) > 0:
                    LOGGER.error(f"Event {event} requires resources {missing_resources} "
                                 f"that are not provided by state {state.get_name()}")

class StateMachineRoot(StateMachine):
    """A State Machine with no parent"""
    def __init__(self, node:Node):
        super().__init__("/", node)

    def new_resource(self, resource:Resource):
        """Add a new resource to the state machine"""
        self.add_required_resource(resource.get_name())
        self._resources[resource.get_name()] = resource.assign(self)

    def set_parent(self, parent):
        LOGGER.error("Cannot set parent of root state machine")
    def _acquire_resource(self, resource_name):
        pass
    def _release_resource(self, resource_name):
        pass
    def _release_all_resources(self):
        pass

class ParallelState(State):
    """A state that runs multiple states in parallel"""
    def __init__(self, name, node:Node = None, states:Iterable[State] = None):
        super().__init__(name, node)
        self._states = []
        if states is not None:
            for state in states:
                self.add_state(state)

    def add_state(self, state):
        """Add a state to the parallel state"""
        self._states.append(state)
        state.set_node(self._node)
        state.set_parent(self._resources)
        self._required_resources |= state.get_required_resources()
        if self._running:
            state.start()

    def start(self):
        """Start the parallel state"""
        super().start()
        for state in self._states:
            state.start()

    def stop(self):
        """Stop the parallel state"""
        for state in self._states:
            state.stop()
        super().stop()

    def step(self):
        """Do one step of the parallel state"""
        super().step()
        for state in self._states:
            state.step()

    def validate(self):
        super().validate()
        # Validate child states
        for state in self._states:
            state.validate()
        # Check resource sharing
        resource_counts = {}
        for state in self._states:
            for resource in state.get_required_resources():
                if resource not in resource_counts:
                    resource_counts[resource] = 0
                resource_counts[resource] += 1
        for resource, count in resource_counts.items():
            if count > 1 and isinstance(resource, UniqueResource):
                LOGGER.warning(f"UniqueResource {resource} is shared by {count} states in"
                               f"parallel state {self._name}. This may cause runtime "
                                "errors or unexpected behavior.")


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from PIL import Image
from matplotlib import pyplot as plt

import numpy as np
import pytesseract
import cv2
import tkinter as tk
import logging
import time
import re
import threading

# Dimensioning values
# We are defining global variables based on match data in order to isolate the scoreboard
left_x = 170
upper_y = 50
right_x = 540
lower_y = 80

time_divide = 230
time_width = 60
time_position = 'left'
# If time-position = right : scoreboard is on the left and time on the right
# Else if time position = left : scoreboard is on the right and time on the left

# To deal with time.sleep() and effectively end the threads
#time_value = 0

class ImageHandler(object):

    def __init__(self, export_path, filename_in):
        self.scoreboard_image = None
        self.time_image = None
        self.time_text = None
        self.teams_goals_image = None
        self.teams_goals_text = None

        self.video_source_path = filename_in
        self.export_image_path = export_path + '/football.jpg'
        self.export_path = export_path

        logging.basicConfig(level=logging.WARNING)

    def extract_image_from_video(self):
        """
        Extracts image from video and saves on disk with specified period.

        :param path_to_video: Path to video and video name with file format
        :param export_image_path: Export image path and image name with file format
        :return: -
        """

        vidcap = cv2.VideoCapture(self.video_source_path)
        count = 0
        #success = True
        image_lst = []

        while(True):
            vidcap.set(cv2.CAP_PROP_POS_MSEC, (count * 1000))
            success, image = vidcap.read()
            image_lst.append(image)

            # Stop when last frame is identified
            if count > 1:
                if np.array_equal(image, image_lst[1]):
                    break
                image_lst.pop(0)  # Clean the list
                # save frame as PNG file

            if(ocr.count < ocr.video_length):
                try:
                    cv2.imwrite(self.export_image_path, image)
                    print('{}.sec reading a new frame: {} '.format(count, success))
                    count += 1
                    ocr.count += 1
                    ocr.eImageExported.set()
                    time.sleep(1)
                except Exception as e:
                    pass

    def localize_scoreboard_image(self):
        """
        Finds the scoreboard table in the upper corner, sets scoreboard_image
        and exports the picture as 'scoreboard_table.jpg'

        :return: True when scoreboard is found
                 False when scoreboard is not found
        """

        # Read a snapshot image from the video and convert to gray
        snapshot_image = cv2.imread(self.export_image_path)
        grayscale_image = cv2.cvtColor(snapshot_image, cv2.COLOR_BGR2GRAY)

        self.scoreboard_image = grayscale_image[upper_y:lower_y,
                                                left_x:right_x]
        cv2.imwrite(self.export_path + '/scoreboard_table.jpg',
                    self.scoreboard_image)

    def split_scoreboard_image(self):
        """
        Splits the scoeboard image into two parts, sets 'time_image' and 'teams_goals_image'
        and exports as 'time_table.jpg' and 'teams_goals_table.jpg'
        Left image represents the time.
        Right image represents the teams and goals.

        :return: -
        """

        '''
        self.time_image = self.scoreboard_image[:, 0:175]
        cv2.imwrite('ocr/img/time_table.jpg', self.time_image)

        self.teams_goals_image = self.scoreboard_image[:, 175:]
        cv2.imwrite('ocr/img/teams_goals_table.jpg', self.teams_goals_image)
        '''

        relative_time_divide = time_divide-left_x
        time_end = relative_time_divide + time_width

        if(time_position == 'right'):
            self.time_image = self.scoreboard_image[:,
                                                    relative_time_divide:time_end]
            cv2.imwrite(self.export_path + '/time_table.jpg', self.time_image)

            self.teams_goals_image = self.scoreboard_image[:,
                                                           0:relative_time_divide]
            cv2.imwrite(self.export_path + '/teams_goals_table.jpg',
                        self.teams_goals_image)

        else:
            self.time_image = self.scoreboard_image[:, 0:relative_time_divide]
            cv2.imwrite(self.export_path + '/time_table.jpg', self.time_image)

            self.teams_goals_image = self.scoreboard_image[:,
                                                           relative_time_divide:]
            cv2.imwrite(self.export_path + '/teams_goals_table.jpg',
                        self.teams_goals_image)

    def enlarge_scoreboard_images(self, enlarge_ratio):
        """
        Enlarges 'time_table.jpg' and 'teams_goals_table.jpg'

        :param enlarge_ratio: Defines the enlarging size (e.g 2-3x)
        :return: -
        """
        self.time_image = cv2.resize(
            self.time_image, (0, 0), fx=enlarge_ratio, fy=enlarge_ratio)
        self.teams_goals_image = cv2.resize(
            self.teams_goals_image, (0, 0), fx=enlarge_ratio, fy=enlarge_ratio)

    def _get_time_from_image(self):
        """
        Preprocesses time_image transformations for OCR.
        Exports 'time_ocr_ready.jpg' after the manipulations.
        Reads match time from 'time_ocr_ready.jpg' using Tesseract.
        Applies result to time_text.

        :return: True: string is found
                 False: string is not found
        """

        # Count nonzero to determine contrast type
        ret, threshed_img = cv2.threshold(
            self.time_image, 0, 255, cv2.THRESH_OTSU+cv2.THRESH_BINARY_INV)

        self.time_image = cv2.GaussianBlur(self.time_image, (3, 3), 0)

        kernel = np.ones((3, 3), np.uint8)

        self.time_image = cv2.erode(self.time_image, kernel, iterations=1)
        self.time_image = cv2.dilate(self.time_image, kernel, iterations=1)

        cv2.imwrite(self.export_path + '/time_ocr_ready.jpg', self.time_image)

        self.time_text = pytesseract.image_to_string(
            Image.open(self.export_path + '/time_ocr_ready.jpg'), config="--psm 6")
        logging.info('Time OCR text: {}'.format(self.time_text))

        if self.time_text is not None:
            return True
        return False

    def _get_teams_goals_from_image(self):
        """
        Preprocesses teams_goals_image with transformations for OCR.
        Exports 'teams_goals_ocr_ready.jpg' after the manipulations.
        Reads teams and goals information from 'teams_goals_ocr_ready.jpg' using Tesseract.
        Applies result to teams_goals_text.

        :return: True: string is found
                 False: string is not found

        """
        # Applying Thresholding for Teams goals OCR preprocess
        ret, self.teams_goals_image = cv2.threshold(
            self.teams_goals_image, 0, 255, cv2.THRESH_OTSU+cv2.THRESH_BINARY_INV)

        self.teams_goals_image = cv2.GaussianBlur(
            self.teams_goals_image, (3, 3), 0)

        kernel = np.ones((3, 3), np.uint8)

        #self.teams_goals_image = cv2.erode(self.teams_goals_image, kernel, iterations=1)
        self.teams_goals_image = cv2.dilate(
            self.teams_goals_image, kernel, iterations=1)

        cv2.imwrite(self.export_path + '/teams_goals_ocr_ready.jpg',
                    self.teams_goals_image)

        self.teams_goals_text = pytesseract.image_to_string(
            Image.open(self.export_path + '/teams_goals_ocr_ready.jpg'))
        logging.info('Teams and goals OCR text: {}'.format(
            self.teams_goals_text))

        if self.teams_goals_text is not None:
            return True
        return False

    def get_scoreboard_texts(self):
        """
        Returns an array of strings including OCR read time, teams and goals texts.
        :return: numpy array 'scoreboard_texts'
                 scoreboard_texts[0] : time text value
                 scoreboard_texts[1] : teams and goals text value

        """
        # Read text values using Tesseract OCR
        time_text_exists = self._get_time_from_image()
        teams_goals_text_exists = self._get_teams_goals_from_image()

        scoreboard_texts = []
        # Use values on successful read
        if time_text_exists and teams_goals_text_exists:
            scoreboard_texts.append(self.time_text)
            scoreboard_texts.append(self.teams_goals_text)
            scoreboard_texts = np.array(scoreboard_texts)

        return scoreboard_texts

    def play_match_video(self):

        cap = cv2.VideoCapture(self.video_source_path)
        count = 0

        if(ocr.time_value < ocr.video_length):
            while (cap.isOpened()):
                cap.set(cv2.CAP_PROP_POS_MSEC, (count * 1000))
                ret, frame = cap.read()

                gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

                cv2.imshow('frame', gray)
                time.sleep(1)
                count += 1
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

            cap.release()
            cv2.destroyAllWindows()

class Match(object):

    def __init__(self, export_path, filename_out):
        self.scoreboard_text_values = None

        self.home_score = 0
        self.home_score_temp = 0

        self.home_team = None
        self.home_team_temp = 0
        self.home_team_fullname = None
        self.home_team_identified = False

        self.opponent_score = 0
        self.opponent_score_temp = 0

        self.opponent_team = None
        self.opponent_team_temp = None
        self.opponent_team_fullname = None
        self.opponent_team_identified = False

        self.match_time = None
        self.match_time_temp = None
        self._match_time_prev = []

        self.export_path = export_path
        self.filename_out = filename_out

    def analize_scoreboard(self):
        while (ocr.count < ocr.video_length):
            try:
                ocr.eImageExported.wait()
                ocr.scoreboard.localize_scoreboard_image()
                ocr.scoreboard.split_scoreboard_image()
                ocr.scoreboard.enlarge_scoreboard_images(3)
                OCR_text = ocr.scoreboard.get_scoreboard_texts()
                ocr.football_match.provide_scoreboard_text_values(OCR_text)
                ocr.football_match.update_all_match_info()
                ocr.football_match.print_all_match_info()
                ocr.eImageExported.clear()
            except Exception as e:
                logging.warning(e)

    def provide_scoreboard_text_values(self, scoreboard_text_values):

        self.scoreboard_text_values = scoreboard_text_values

    def cleanse_match_score(self):
        """
        Cleanse home_score_temp and opponent_score_temp values and removes
        noisy starters and enders if present

        :return: -
        """
        score_string = self.scoreboard_text_values[1].split(' ')[1]

        result = []
        for letter in score_string:
            if letter.isdigit():
                result += letter
        self.home_score_temp = result[0]
        self.opponent_score_temp = result[1]

    def cleanse_match_teams(self):
        """
        Cleanse home_team_temp and opponent_team_temp values and removes
        noisy starter or ender if present

        :return: -
        """
        self.home_team_temp = self.scoreboard_text_values[1].split(' ')[0]
        self.opponent_team_temp = self.scoreboard_text_values[1].split(' ')[2]

        # Check and remove noisy starters and enders
        if not self.home_team_temp[0].isalpha():
            self.home_team_temp = self.home_team_temp[1:4]
        elif not self.opponent_team_temp[-1].isalpha():
            self.opponent_team_temp = self.opponent_team_temp[0:3]

    def cleanse_match_time(self):
        """
        Cleanse match_time_temp, and removes noisy starter or ender if present

        :return: -
        """

        self.match_time_temp = self.scoreboard_text_values[0]

        # Check for noisy starters and ender and clean if present
        letter_ptr = 0
        if not self.match_time_temp[letter_ptr].isdigit():
            letter_ptr += 1
        if not self.match_time_temp[letter_ptr].isdigit():
            letter_ptr += 1
            self.match_time_temp = self.match_time_temp[letter_ptr:]
            logging.info("Time text noisy starter removed.")
        elif not self.match_time_temp[-1].isdigit():
            self.match_time_temp = self.match_time_temp[0:-1]
            logging.info("Time text noisy ender removed.")

    def update_match_time(self):
        """
        Validates cleansed match_time_temp with regular expression and sets match_time if valid value exists

        :return: True: time has been updated
                 False: time has not been updated
        """

        # Check if the OCR read value is valid
        time_expr = re.compile('\d\d:\d\d')
        res = time_expr.search(self.match_time_temp)

        if res is None:
            return False

        last_valid_timeval = self.match_time_temp[res.start():res.end()]
        self._match_time_prev.append(last_valid_timeval)

        # Check validity between last time values
        if last_valid_timeval < self._match_time_prev[len(self._match_time_prev)-2]:
            # Minute error occured - minute remain unchanged
            if last_valid_timeval[0:2] < self._match_time_prev[len(self._match_time_prev)-2][0:2]:
                logging.warning(
                    "Minute error occured: minute remain unchanged!")
                fixed_minutes = self._match_time_prev[len(
                    self._match_time_prev)-2][0:2]
                last_valid_timeval = fixed_minutes + last_valid_timeval[2:]

            else:
                # Second error occured - auto increment second
                logging.warning(
                    "Second error occured: auto incremented second!")
                seconds = self._match_time_prev[len(
                    self._match_time_prev)-2][-2:]
                fixed_seconds = str(int(seconds)+1)
                last_valid_timeval = last_valid_timeval[:-2] + fixed_seconds

        # Free unnecessary time values
        if len(self._match_time_prev) > 2:
            self._match_time_prev.pop(0)

        # Write all valid values to a text file for analysis
        self.match_time = last_valid_timeval
        with open(self.export_path + '/' + self.filename_out, 'a') as f:
            f.write("%s,%s\n" % (self.match_time, ocr.count))
        return True

    def update_match_score(self):
        """
        Validates cleansed score with regular expression

        :return: True: score matches the regexp
                 False: score does not match the regexp
        """
        score_expr = re.compile('\d-\d')
        res = score_expr.search(self.scoreboard_text_values[1])

        if res is None:
            return False

        self.home_score = self.home_score_temp
        self.opponent_score = self.opponent_score_temp
        return True

    def update_match_team(self, selected_team):
        """
        Sets cleansed home_team or opponent_team values if not set before

        :return: -
        """
        if selected_team == 'home':
            self.home_team = self.home_team_temp
            self.home_team_identified = True

        elif selected_team == 'opponent':
            self.opponent_team = self.opponent_team_temp
            self.opponent_team_identified = True

    def update_all_match_info(self):
        """
        Attempts to update match infos:
        time, teams, score
        :return: True: update succeed
                 False: update failed
        """
        if len(self.scoreboard_text_values[0]) > 0 and len(self.scoreboard_text_values[1]) > 0:
            try:
                # Clean OCR read time value and update time if valid
                self.cleanse_match_time()
                self.update_match_time()

                # Clean OCR read score value and update score if valid
                self.cleanse_match_score()
                self.update_match_score()

                # Clean OCR read team values and set teams if valid and necessary
                self.cleanse_match_teams()

                if self.home_team_identified is False:
                    self.update_match_team('home')

                if self.opponent_team_identified is False:
                    self.update_match_team('opponent')

            except Exception as e:
                logging.info(e)
                logging.info("Unable to update match info for some reason")
        else:
            logging.info("Unable to update match info: no text received!")

    def print_all_match_info(self):

        home_team_name = self.home_team
        opponent_team_name = self.opponent_team

        if self.home_team_fullname is not None and self.opponent_team_fullname is not None:
            home_team_name = self.home_team_fullname
            opponent_team_name = self.opponent_team_fullname

        print('{} {} {}-{} {}'.format(self.match_time,
                                      home_team_name,
                                      self.home_score,
                                      self.opponent_score,
                                      opponent_team_name))

# MAIN
# Empty times.txt file
def ocr(export_path, filename_in, filename_out, video_length):
    ocr.count = 0
    ocr.video_length = video_length
    open(export_path+'/' + filename_out, 'w').close()

    ocr.eImageExported = threading.Event()

        # Create objects and threads
    ocr.scoreboard = ImageHandler(export_path, filename_in)
    ocr.football_match = Match(export_path, filename_out)

    ocr.tImageExtractor = threading.Thread(
            None, ocr.scoreboard.extract_image_from_video, name="ImageExtractor")
    ocr.tScoreboardAnalyzer = threading.Thread(
            None, ocr.football_match.analize_scoreboard, name="ScoreboardAnalyzer")

    ocr.tImageExtractor.start()
    ocr.tScoreboardAnalyzer.start()

    ocr.tImageExtractor.join()
    ocr.tScoreboardAnalyzer.join()

if __name__ == '__main__' :
    filename_in = 'ocr/tmp/secondmatch.mkv'
    export_path = 'ocr/img'
    filename_out = 'times.txt'
    video_length = 1080
    ocr(export_path, filename_in, filename_out, 1080)


import torch
import torch.nn as nn

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        self.fc1 = nn.Linear(784, 256, bias = False)
        self.bn1 = nn.BatchNorm1d(256)
        self.relu1 = nn.ReLU()

        self.fc2 = nn.Linear(256, 128)
        self.bn2 = nn.BatchNorm1d(128)
        self.relu2 = nn.ReLU()

        self.fc3 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = self.bn1(x)
        x = self.relu1(x)

        x = self.fc2(x)
        x = self.bn2(x)
        x = self.relu2(x)

        x = self.fc3(x)
        return x


import torch
import torch.nn as nn
from math import sin, cos
import models
from models.base import BaseModel
from models.utils import chunk_batch
from systems.utils import update_module_step
from nerfacc import ContractionType, OccupancyGrid, ray_marching
from nerfacc.vol_rendering import render_transmittance_from_alpha, rendering
from utils.rotation import R_from_quaternions

@models.register('se3')
class SE3Model(BaseModel):
    def setup(self):
        self.static_geometry = models.make(self.config.geometry.name, self.config.geometry)
        self.static_texture = models.make(self.config.texture.name, self.config.texture)
        self.dynamic_geometry = models.make(self.config.geometry.name, self.config.geometry)
        self.dynamic_texture = models.make(self.config.texture.name, self.config.texture)

        init_angle = self.config.get('init_angle', 0.1)
        init_dir = self.config.get('init_dir', [1., 1., 1.])
        self.quaternions = nn.Parameter(self.init_quaternions(half_angle=init_angle, init_dir=init_dir), requires_grad=True)
        self.translation = nn.Parameter(torch.tensor([0.001, 0.001, 0.001], dtype=torch.float32), requires_grad=True)
        self.canonical = 0.5

        self.register_buffer('scene_aabb', torch.as_tensor([-self.config.radius, -self.config.radius, -self.config.radius, self.config.radius, self.config.radius, self.config.radius], dtype=torch.float32))
        if self.config.grid_prune:
            self.grid_warmup = self.config['grid_warmup']
            self.occupancy_grid = OccupancyGrid(
                roi_aabb=self.scene_aabb,
                resolution=128, # the resolution is open to discuss
                contraction_type=ContractionType.AABB
            )

        self.randomized = self.config.randomized
        if self.config.white_bkgd:
            self.register_buffer('background_color', torch.as_tensor([1.0, 1.0, 1.0], dtype=torch.float32), persistent=False)
            self.background_color.to(self.rank)

        self.render_step_size = 1.732 * 2 * self.config.radius / self.config.num_samples_per_ray

    def update_step(self, epoch, global_step):
        update_module_step(self.static_texture, epoch, global_step)
        update_module_step(self.dynamic_texture, epoch, global_step)

        def occ_eval_fn(x):
            density_s, _ = self.static_geometry(x)
            x_d = self.rigid_transform(x)
            density_d, _ = self.dynamic_geometry(x_d)
            density = density_s + density_d
            return density[...,None] * self.render_step_size

        if self.training and self.config.grid_prune:
            self.occupancy_grid.every_n_step(step=global_step, occ_eval_fn=occ_eval_fn, occ_thre=1e-4, warmup_steps=self.grid_warmup)

    def isosurface(self):
        mesh_s = self.static_geometry.isosurface()
        mesh_d = self.dynamic_geometry.isosurface()
        return {'static': mesh_s, 'dynamic': mesh_d}

    def init_quaternions(self, half_angle, init_dir):
        a = torch.tensor([init_dir[0], init_dir[1], init_dir[2]], dtype=torch.float32)
        a = torch.nn.functional.normalize(a, p=2., dim=0)
        sin_ = sin(half_angle)
        cos_ = cos(half_angle)
        r = cos_
        i = a[0] * sin_
        j = a[1] * sin_
        k = a[2] * sin_
        q = torch.tensor([r, i, j, k], dtype=torch.float32)
        return q

    def rigid_transform(self, positions, state=0.):
        '''
        Perform the rigid transformation: R_axis_d,rot_angle(center=axis_o) @ x + t
        '''

        scaling = (self.canonical - state) / self.canonical

        if scaling == 1.:
            R = R_from_quaternions(self.quaternions)
            positions = torch.matmul(R, positions.T).T
            positions = positions + self.translation
        elif scaling == -1.:
            positions = positions - self.translation
            inv_sc = torch.tensor([1., -1., -1., -1]).to(self.quaternions)
            inv_q = inv_sc * self.quaternions
            R = R_from_quaternions(inv_q)
            positions = torch.matmul(R, positions.T).T
        else:
            raise NotImplementedError
        return positions

    def forward_(self, rays, scene_state):
        rays_o, rays_d = rays[:, 0:3], rays[:, 3:6] # both (N_rays, 3)

        def sigma_fn_composite(t_starts, t_ends, ray_indices):
            ray_indices = ray_indices.long()
            t_origins = rays_o[ray_indices]
            t_dirs = rays_d[ray_indices]
            positions = t_origins + t_dirs * (t_starts + t_ends) / 2.
            sigma_s, _ = self.static_geometry(positions)
            positions = self.rigid_transform(positions, scene_state)
            sigma_d, _ = self.dynamic_geometry(positions)
            sigma = sigma_s + sigma_d
            return sigma[...,None]

        def rgb_sigma_fn_static(t_starts, t_ends, ray_indices):
            ray_indices = ray_indices.long()
            t_origins = rays_o[ray_indices]
            t_dirs = rays_d[ray_indices]
            positions = t_origins + t_dirs * (t_starts + t_ends) / 2.
            density, feature = self.static_geometry(positions)
            rgb = self.static_texture(feature, t_dirs)
            return rgb, density[...,None]

        def rgb_sigma_fn_dynamic(t_starts, t_ends, ray_indices):
            ray_indices = ray_indices.long()
            t_origins = rays_o[ray_indices]
            t_dirs = rays_d[ray_indices]
            positions = t_origins + t_dirs * (t_starts + t_ends) / 2.
            positions = self.rigid_transform(positions, scene_state)
            density, feature = self.dynamic_geometry(positions)
            dirs_d = self.rigid_transform(t_dirs, scene_state)
            rgb = self.dynamic_texture(feature, dirs_d)
            return rgb, density[...,None]

        def composite_rendering(ray_indices, t_starts, t_ends):
            n_rays = rays_o.shape[0]
            rgb_s, sigma_s = rgb_sigma_fn_static(t_starts, t_ends, ray_indices)
            rgb_d, sigma_d = rgb_sigma_fn_dynamic(t_starts, t_ends, ray_indices)

            dists = t_ends - t_starts
            alpha_s = 1. - torch.exp(-sigma_s * dists)
            alpha_d = 1. - torch.exp(-sigma_d * dists)

            alpha_add = 1. - (1. - alpha_s) * (1. - alpha_d)
            Ts = render_transmittance_from_alpha(alpha_add, ray_indices=ray_indices)

            weights_s = alpha_s * Ts
            weights_d = alpha_d * Ts

            weights = weights_s + weights_d
            # opacity
            opacity = self.acc_along_rays(weights, ray_indices, n_rays)
            opacity = opacity.squeeze(-1)
            # acc color
            rgb = weights_s * rgb_s + weights_d * rgb_d
            rgb = self.acc_along_rays(rgb, ray_indices, n_rays)
            # Background composition.
            if self.config.white_bkgd:
                rgb = rgb + self.background_color * (1. - opacity[..., None])

            # validation and testing
            if not self.training:
                # depth
                depth = weights * ((t_starts + t_ends) * 0.5)
                depth = self.acc_along_rays(depth, ray_indices, n_rays)
                depth = depth.squeeze(-1)

                rgb_s_only, opacity_s, depth_s_only = rendering(t_starts, t_ends, ray_indices, n_rays,
                                                                rgb_sigma_fn=rgb_sigma_fn_static,
                                                                render_bkgd=self.background_color)
                rgb_d_only, opacity_d, depth_d_only = rendering(t_starts, t_ends, ray_indices, n_rays,
                                                                rgb_sigma_fn=rgb_sigma_fn_dynamic,
                                                                render_bkgd=self.background_color)
                return {
                    'rgb': rgb,
                    'opacity': opacity,
                    'depth': depth,
                    'rgb_s': rgb_s_only,
                    'rgb_d': rgb_d_only,
                    'depth_s': depth_s_only,
                    'depth_d': depth_d_only,
                    'opacity_s': opacity_s,
                    'opacity_d': opacity_d,
                }

            return {
                'rgb': rgb,
                'rays_valid': opacity > 0,
                'opacity': opacity,
            }

        with torch.no_grad():
            ray_indices, t_starts, t_ends = ray_marching(
                rays_o, rays_d,
                scene_aabb=self.scene_aabb,
                grid=self.occupancy_grid if self.config.grid_prune else None,
                sigma_fn=sigma_fn_composite,
                render_step_size=self.render_step_size,
                stratified=self.randomized,
            )
        render_out = composite_rendering(ray_indices, t_starts, t_ends)

        if self.training:
            return {
                'comp_rgb': render_out['rgb'],
                'opacity': render_out['opacity'],
                'rays_valid': render_out['rays_valid'],
                'num_samples': torch.as_tensor([len(t_starts)], dtype=torch.int32, device=rays.device),
            }

        return {
            'comp_rgb': render_out['rgb'],
            'opacity': render_out['opacity'],
            'depth': render_out['depth'],
            'rgb_s': render_out['rgb_s'],
            'rgb_d': render_out['rgb_d'],
            'depth_s': render_out['depth_s'],
            'depth_d': render_out['depth_d'],
            'opacity_s': render_out['opacity_s'],
            'opacity_d': render_out['opacity_d'],
        }

    def forward(self, rays_0, rays_1):
        if self.training:
            out_0 = self.forward_(rays_0, scene_state=0.)
            out_1 = self.forward_(rays_1, scene_state=1.)
        else:
            out_0 = chunk_batch(self.forward_, self.config.ray_chunk, rays_0, scene_state=0.)
            out_1 = chunk_batch(self.forward_, self.config.ray_chunk, rays_1, scene_state=1.)
            del rays_0, rays_1
        return [{**out_0}, {**out_1}]

    def train(self, mode=True):
        self.randomized = mode and self.config.randomized
        return super().train(mode=mode)

    def eval(self):
        self.randomized = False
        return super().eval()

    def regularizations(self, outs):
        losses = {}
        return losses


import numpy as np
import pandas as pd
from math import factorial, pi
import scipy.optimize
import scipy.misc
import os
import re
import argparse
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF,ConstantKernel
# for tests
#import matplotlib.pyplot as pl

# directory where the fit_lattice_test.py is located
dir_path = os.path.dirname(os.path.realpath(__file__))

########################################################################
if __name__ == "__main__":
    __doc__="""Construct a parametrization (from PhysRevC.100.064910) of the lattice QCD equation of state
(P/T^4, n/T^3, s/T^3, e/T^4) by calling function:
- param(T,muB,muQ,muS)
  input: temperature and chemical potentials in [GeV]
  output: dictionnary of all quantities ['T','P','s','n_B','n_Q','n_S','e']

Produces lattice data for P/T^4, nB/T^3, s/T^3, e/T^4 as a function of T for a single value of muB:
- lattice_data(EoS,muB)
  input: - EoS: - 'muB' refers to the EoS with the condition \mu_Q = \mu_S = 0
                - 'nS0' refers to the EoS with the condition <n_S> = 0 & <n_Q> = 0.4 <n_B>
         - muB: baryon chemical potential in [GeV]
  output: dictionnary of all quantities + error ['T','P','s','n_B','e']

Calculation of the equation of state under the conditions: <n_S> = 0 ; <n_Q> = factQB*<n_B>:
- EoS_nS0(fun,T,muB,**kwargs)
  input: - fun: any function which calculate an EoS (by ex: param, HRG, full_EoS)
         - T,muB: temperature and baryon chemical potential in [GeV]
  output: dictionnary of all quantities ['T','P','s','n_B','e']
"""
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter
    )
    args = parser.parse_args()

###############################################################################
# J. Phys.: Conf. Ser. 1602 012011
# critical temperature from lattice at \mu_B = 0
Tc0 = 0.158
# expansion coefficients of T_c(\mu_B)
kappa2 = 0.0153
kappa4 = 0.00032
###############################################################################
def Tc_lattice(muB):
    """
    Critical temperature as a function of muB from lQCD
    J. Phys.: Conf. Ser. 1602 012011
    """
    return Tc0*(1.-kappa2*(muB/Tc0)**2.-kappa4*(muB/Tc0)**4.)

###############################################################################
def dTcdmuB_lattice(muB):
    """
    Derivative of the critical temperature wrt \mu_B
    """
    dTc = -2.*muB*kappa2/Tc0 -4.*(muB**3.)*kappa4/Tc0**3.
    return dTc

def Tc_lattice_muBoT(muBoT):
    """
    Find the critical temperature Tc for a fixed muB/T
    """
    if(muBoT==0):
        return Tc_lattice(0.)
    else:
        xmuB = scipy.optimize.root(lambda muB: muB/Tc_lattice(muB)-muBoT,[muBoT*Tc0],method='lm').x[0]
        return Tc_lattice(xmuB)

###############################################################################
# SB limits
###############################################################################
def SB_lim(T,muB,muQ,muS,Nf=3):
    """
    SB limit for p,s,n,e,cs^2
    """
    Nc = 3 # number of colors
    dgq = 2.*Nc # degeneracy factor for quarks
    dgg = 2.*(Nc**2.-1.) # degeneracy factor for gluons

    # if input is a single temperature value T
    if(isinstance(T,float)):

        # chemical potential asssociated to each quark
        mu_u = 1./3.*muB + 2./3.*muQ
        mu_d = 1./3.*muB - 1./3.*muQ
        mu_s = 1./3.*muB - 1./3.*muQ - muS

        # list of partons corresponding to the given number of flavors
        # chemical potentials and charges
        if(Nf==0):
            list_mu = []
            list_B = []
            list_Q = []
            list_S =[]
        elif(Nf==2):
            list_mu = [mu_u,mu_d]
            list_B = [1./3.,1./3.]
            list_Q = [2./3.,-1./3.]
            list_S =[0.,0.]
        elif(Nf==3):
            list_mu = [mu_u,mu_d,mu_s]
            list_B = [1./3.,1./3.,1./3.]
            list_Q = [2./3.,-1./3.,-1./3.]
            list_S = [0.,0.,-1.]

        Pgluons = (pi**2)/90.*dgg
        Pquarks = dgq*sum([(7*pi**2)/360. + ((mu_q/T)**2)/12. + ((mu_q/T)**4)/(24*pi**2) for mu_q in list_mu])
        P = Pgluons+Pquarks

        sgluons = 4*(pi**2)/90.*dgg
        squarks = dgq*sum([(4*7*pi**2)/360. + 2*((mu_q/T)**2)/12.  for mu_q in list_mu])
        s = sgluons+squarks

        nB = dgq*sum([B_q*((2*(mu_q/T))/12. + (4*(mu_q/T)**3)/(24*pi**2)) for B_q,mu_q in zip(list_B,list_mu)])
        nQ = dgq*sum([Q_q*((2*(mu_q/T))/12. + (4*(mu_q/T)**3)/(24*pi**2)) for Q_q,mu_q in zip(list_Q,list_mu)])
        nS = dgq*sum([S_q*((2*(mu_q/T))/12. + (4*(mu_q/T)**3)/(24*pi**2)) for S_q,mu_q in zip(list_S,list_mu)])

        e = s-P+muB/T*nB+muQ/T*nQ+muS/T*nS

        cs2 = 1./3.

    # if the input is a list of temperature values
    elif(isinstance(T,np.ndarray) or isinstance(T,list)):

        P = np.zeros_like(T)
        s = np.zeros_like(T)
        nB = np.zeros_like(T)
        nQ = np.zeros_like(T)
        nS = np.zeros_like(T)
        e = np.zeros_like(T)
        cs2 = 1./3.*np.ones_like(T)
        for i,xT in enumerate(T):
            # see if arrays are also given for chemical potentials
            try:
                xmuB = muB[i]
            except:
                xmuB = muB
            try:
                xmuQ = muQ[i]
            except:
                xmuQ = muQ
            try:
                xmuS = muS[i]
            except:
                xmuS = muS

            result = SB_lim(xT,xmuB,xmuQ,xmuS,Nf=Nf)

            P[i] = result['P']
            s[i] = result['s']
            nB[i] = result['n_B']
            nQ[i] = result['n_Q']
            nS[i] = result['n_S']
            e[i] = result['e']

    else:
        raise Exception('Problem with input')

    return {'P':P,'s':s,'n_B':nB,'n_Q':nQ,'n_S':nS,'e':e,'I':e-3.*P,'cs^2':cs2}

###############################################################################
# import data for the parametrization of susceptibilities
###############################################################################

chi_a = {}
for chi_file in ["/data/chi_a_nS0.csv","/data/chi_a.csv"]:
    param_chi_a = pd.read_csv(dir_path+chi_file).to_dict(orient='list')
    # scan rows for each chi
    for i,chi in enumerate(param_chi_a['chi']):
        values = []
        # scan columns with coefficients
        for j,coeff in enumerate(param_chi_a):
            # skip first column which is chi string
            if(coeff=='chi'):
                continue
            # append values
            values.append(param_chi_a[coeff][i])
        chi_a.update({chi:values})

chi_b = {}
for chi_file in ["/data/chi_b_nS0.csv","/data/chi_b.csv"]:
    param_chi_b = pd.read_csv(dir_path+chi_file).to_dict(orient='list')
    # scan rows for each chi
    for i,chi in enumerate(param_chi_b['chi']):
        values = []
        # scan columns with coefficients
        for j,coeff in enumerate(param_chi_b):
            # skip first column which is chi string
            if(coeff=='chi'):
                continue
            # append values
            values.append(param_chi_b[coeff][i])
        chi_b.update({chi:values})

# list of all susceptibilities
list_chi = list(param_chi_a['chi'])
list_chi_nS0 = ['chiB2_nS0','chiB4_nS0']

########################################################################
# Stefan Boltzmann limit for the susceptibilities
# can be found in PhysRevC.100.064910
chi_SB = dict(zip(list_chi,[19.*pi**2./36.,\
    1./3.,2./3.,1.,\
    0.,-1./3.,1./3.,\
    2./(9.*pi**2.),4./(3*pi**2.),6./pi**2.,\
    0.,-2./(9.*pi**2.),2./(9.*pi**2.),\
    4./(9.*pi**2.),-2./pi**2.,2./pi**2.,\
    4./(9.*pi**2.),2./(3.*pi**2.),2./(3.*pi**2.),\
    2./(9.*pi**2.),-2./(9.*pi**2.),-2./(3.*pi**2.)]))

chi_SB.update(dict(zip(list_chi_nS0,[0.1067856506125367,0.0006673764465596013])))

########################################################################
def param_chi(T,quant):
    """
    Parametriation of the susceptibilities at as a function of temperature
    Ex: param_chi(T,'chiBQS121')
    input quant is a string with the format: chiBQS121
    input T being a list or a float
    """
    tt = T/Tc_lattice(0.)
    numerator = sum([ai/(tt)**i for i,ai in enumerate(chi_a[quant])])
    denominator = sum([bi/(tt)**i for i,bi in enumerate(chi_b[quant])])
    c0 = chi_SB[quant]-chi_a[quant][0]/chi_b[quant][0]
    return numerator/denominator + c0

########################################################################
# for each susceptibility, get the order of the derivative wrt B,Q,S
########################################################################

BQS = dict(zip(list_chi,[{'B': 0, 'Q': 0, 'S': 0} for i in range(len(list_chi))]))
chi_latex = {'chi0':r'$\chi_0$'}
for chi in list_chi:
    # derivatives wrt to each charge
    if(chi!='chi0'):
        # decompose chiBQS234 in [B,Q,S] and [2,3,4]
        chi_match = re.match('chi([A-Z]+)([0-9]+)', chi)
        list_charge = list(chi_match.group(1)) # contains the charges
        list_der = list(chi_match.group(2)) # contains the derivatives
        chi_latex.update({chi:r'$\chi^{'+"".join(list_charge)+'}_{'+"".join(list_der)+'}$'})
        for ich,xcharge in enumerate(list_charge):
            BQS[chi][xcharge] = int(list_der[ich]) # match each charge to its derivative

chi_latex.update({'chiB2_nS0':r'$c_2$', 'chiB4_nS0':r'$c_4$'})

########################################################################
def param(T,muB,muQ,muS):
    """
    Parametrization of thermodynamic quantities from lQCD
    as a function of T, \mu_B, \mu_Q, \mu_S
    """
    # if input is a single temperature value T
    if(isinstance(T,float)):
        p = 0.
        nB = 0.
        nQ = 0.
        nS = 0.
        s = 0.
        e = 0.

        if(muB==0. and muQ==0. and muS==0.):
            p = param_chi(T,'chi0')
            der = scipy.misc.derivative(param_chi,T,dx=1e-5,args=('chi0',))
            s = T*der
        else:
            for chi in list_chi:
                i = BQS[chi]['B']
                j = BQS[chi]['Q']
                k = BQS[chi]['S']
                fact = 1./(factorial(i)*factorial(j)*factorial(k))
                xchi = param_chi(T,chi)
                pow_muB = ((muB/T)**i)
                pow_muQ = ((muQ/T)**j)
                pow_muS = ((muS/T)**k)
                # pressure P/T^4
                p += fact*xchi*pow_muB*pow_muQ*pow_muS
                # baryon density n_B/T^3 when i > 1
                if(i >= 1):
                    nB += fact*xchi*i*((muB/T)**(i-1.))*pow_muQ*pow_muS
                # charge density n_Q/T^3 when i > 1
                if(j >= 1):
                    nQ += fact*xchi*pow_muB*j*((muQ/T)**(j-1.))*pow_muS
                # strangeness density n_S/T^3 when k > 1
                if(k >= 1):
                    nS += fact*xchi*pow_muB*pow_muQ*k*((muS/T)**(k-1.))
                # derivative of the susceptibility wrt temperature
                der = scipy.misc.derivative(param_chi,T,dx=1e-5,args=(chi,))
                # s/T^3 = T d(P/T^4)/d(T) + 4 P/T^4
                # here we add just the 1st part
                s += fact*(T*der-(i+j+k)*xchi)*pow_muB*pow_muQ*pow_muS
        # add 2nd piece to s/T^3
        s += 4.*p
        # energy density e/T^4
        e = s-p+(muB/T)*nB+(muQ/T)*nQ+(muS/T)*nS

    # if the input is a list of temperature values
    elif(isinstance(T,np.ndarray) or isinstance(T,list)):
        p = np.zeros_like(T)
        s = np.zeros_like(T)
        nB = np.zeros_like(T)
        nQ = np.zeros_like(T)
        nS = np.zeros_like(T)
        e = np.zeros_like(T)
        for i,xT in enumerate(T):
            # see if arrays are also given for chemical potentials
            try:
                xmuB = muB[i]
            except:
                xmuB = muB
            try:
                xmuQ = muQ[i]
            except:
                xmuQ = muQ
            try:
                xmuS = muS[i]
            except:
                xmuS = muS
            result = param(xT,xmuB,xmuQ,xmuS)
            p[i] = result['P']
            s[i] = result['s']
            nB[i] = result['n_B']
            nQ[i] = result['n_Q']
            nS[i] = result['n_S']
            e[i] = result['e']

    else:
        raise Exception('Problem with input')

    return {'T': T,'P':p, 's':s, 'n_B':nB, 'n_Q':nQ, 'n_S':nS, 'e':e, 'I':e-3*p}

########################################################################
def param_nS0(T,muB):
    """
    Parametrization of thermodynamic quantities from lQCD
    as a function of T, \mu_B for the case <n_S>=0 & <n_Q>=0.4<n_B>
    """
    # if input is a single temperature value T
    if(isinstance(T,float)):
        p = 0.
        nB = 0.
        nQ = 0.
        nS = 0.
        s = 0.
        e = 0.

        p = param_chi(T,'chi0')
        der = scipy.misc.derivative(param_chi,T,dx=1e-5,args=('chi0',))
        s = T*der
        if(muB!=0.):
            for ichi,chi in enumerate(list_chi_nS0):
                i = 2*(ichi+1)
                xchi = param_chi(T,chi)
                pow_muB = ((muB/T)**i)
                # pressure P/T^4
                p += xchi*pow_muB
                # baryon density n_B/T^3 when i > 1
                nB += xchi*i*((muB/T)**(i-1.))
                # derivative of the susceptibility wrt temperature
                der = scipy.misc.derivative(param_chi,T,dx=1e-5,args=(chi,))
                # s/T^3 = T d(P/T^4)/d(T) + 4 P/T^4
                # here we add just the 1st part
                s += (T*der-(i)*xchi)*pow_muB
        # add 2nd piece to s/T^3
        s += 4.*p
        # energy density e/T^4
        e = s-p+(muB/T)*nB
        # charge density
        nQ = 0.4*nB

    # if the input is a list of temperature values
    elif(isinstance(T,np.ndarray) or isinstance(T,list)):
        p = np.zeros_like(T)
        s = np.zeros_like(T)
        nB = np.zeros_like(T)
        nQ = np.zeros_like(T)
        nS = np.zeros_like(T)
        e = np.zeros_like(T)
        for i,xT in enumerate(T):
            # see if arrays are also given for chemical potentials
            try:
                xmuB = muB[i]
            except:
                xmuB = muB
            result = param_nS0(xT,xmuB)
            p[i] = result['P']
            s[i] = result['s']
            nB[i] = result['n_B']
            nQ[i] = result['n_Q']
            nS[i] = result['n_S']
            e[i] = result['e']

    else:
        raise Exception('Problem with input')

    return {'T': T,'P':p, 's':s, 'n_B':nB, 'n_Q':nQ, 'n_S':nS, 'e':e, 'I':e-3*p}

###############################################################################
# import data from lattice at muB = 0
###############################################################################

# read chi0
WB_EoS0 = pd.read_csv(dir_path+"/data/WB-EoS_muB0_j.physletb.2014.01.007.csv").to_dict(orient='list')
chi_lattice2014 = {'chi0':np.array(list(zip(WB_EoS0['T'],WB_EoS0['P'],WB_EoS0['P_err'])))}
# save all other thermodynamic quantities
for quant in WB_EoS0:
    WB_EoS0[quant] = np.array(WB_EoS0[quant])
# read data from 2012 (chiB2,chiQ2,chiS2)
chi_lattice2012 = {}
try:
    df = pd.read_csv(dir_path+"/data/WB_chi_T_JHEP01(2012)138.csv").to_dict(orient='list')
    for entry in df:
        if(entry=='T' or '_err' in entry):
            continue
        chi_lattice2012.update({entry:np.array(list(zip(df['T'],df[entry],df[entry+'_err'])))})
except:
    pass
# read data from 2015 (chiB2,chiB4,chiS2)
chi_lattice2015 = {}
try:
    df = pd.read_csv(dir_path+"/data/WB_chi_T_PhysRevD.92.114505.csv").to_dict(orient='list')
    for entry in df:
        if(entry=='T' or '_err' in entry):
            continue
        chi_lattice2015.update({entry:np.array([[df['T'][iT],df[entry][iT],df[entry+'_err'][iT]] for iT,_ in enumerate(df[entry]) if np.logical_not(np.isnan(df[entry][iT]))])})
except:
    pass
# read data from 2017 (chiB2,chiB4,chiB2) for <nS>=0 & <nQ>=0.4<nB>
chi_lattice2017 = {}
try:
    df = pd.read_csv(dir_path+"/data/WB_chi_nS0_T_EPJWebConf.137(2017)07008.csv").to_dict(orient='list')
    for entry in df:
        if(entry=='T' or '_err' in entry):
            continue
        chi_lattice2017.update({entry:np.array([[df['T'][iT],df[entry][iT],df[entry+'_err'][iT]] for iT,_ in enumerate(df[entry]) if np.logical_not(np.isnan(df[entry][iT]))])})
except:
    pass
# read data from 2018
chi_lattice2018 = {}
try:
    df = pd.read_csv(dir_path+"/data/WB_chi_T_JHEP10(2018)205.csv").to_dict(orient='list')
    for entry in df:
        if(entry=='T' or '_err' in entry):
            continue
        chi_lattice2018.update({entry:np.array(list(zip(df['T'],df[entry],df[entry+'_err'])))})
except:
    pass
# read data from 2020 (chiBQ11,chiBS11,chiQS11)
chi_lattice2020 = {}
try:
    df = pd.read_csv(dir_path+"/data/WB_chi_T_PhysRevD.101.034506.csv").to_dict(orient='list')
    for entry in df:
        if(entry=='T' or '_err' in entry):
            continue
        chi_lattice2020.update({entry:np.array(list(zip(df['T'],df[entry],df[entry+'_err'])))})
except:
    pass
# read data from 2021
WB_EoS_muBoT2021 = {}
try:
    df = pd.read_csv(dir_path+"/data/WB-EoS_muBoT_2102.06660.csv").to_dict(orient='list')
    for entry in df:
        if(entry=='T' or '_err' in entry):
            continue
        WB_EoS_muBoT2021.update({entry:np.array(list(zip(df['T'],df[entry],df[entry+'_err'])))})
except:
    pass

###############################################################################
def EoS_nS0(fun,T,muB,**kwargs):
    """
    Calculation of the EoS defined by the input function at (T,muB) with the conditions:
    <n_S> = 0
    <n_Q> = factQB*<n_B>
    """
    factQB = 0.4

    if(isinstance(T,float)):
        p = 0.
        nB = 0.
        nQ = 0.
        nS = 0.
        s = 0.
        e = 0.
        n = 0.
        chi = np.zeros(len(list_chi))

        def system(mu):
            """
            Define the system to be solved
            <n_S> = 0
            <n_Q> = factQB * <n_B>
            """
            thermo = fun(T,muB,mu[0],mu[1],**kwargs)
            return [thermo['n_S']*T**3, thermo['n_Q']*T**3-factQB*thermo['n_B']*T**3]

        solution = scipy.optimize.root(system,[-0.08*muB,0.03*muB],method='lm').x
        muQ = solution[0]
        muS = solution[1]
        result = fun(T,muB,muQ,muS,**kwargs)

        p = result['P']
        s = result['s']
        nB = result['n_B']
        nQ = factQB*nB
        nS = 0.
        e = result['e']
        # some extra quantities are calculated within HRG function
        try:
            n = result['n']
            chi = result['chi']
        except:
            pass

    elif(isinstance(T,np.ndarray) or isinstance(T,list)):
        p = np.zeros_like(T)
        s = np.zeros_like(T)
        nB = np.zeros_like(T)
        nQ = np.zeros_like(T)
        nS = np.zeros_like(T)
        n = np.zeros_like(T)
        e = np.zeros_like(T)
        muQ = np.zeros_like(T)
        muS = np.zeros_like(T)
        chi = np.zeros((len(list_chi),len(T)))
        for i,xT in enumerate(T):
            # see if arrays are also given for chemical potentials
            try:
                xmuB = muB[i]
            except:
                xmuB = muB
            result = EoS_nS0(fun,xT,xmuB,**kwargs)
            p[i] = result['P']
            s[i] = result['s']
            nB[i] = result['n_B']
            nQ[i] = result['n_Q']
            nS[i] = result['n_S']
            n[i] = result['n']
            e[i] = result['e']
            muQ[i] = result['muQ']
            muS[i] = result['muS']
            chi[:,i] = result['chi']

    else:
        raise Exception('Problem with input')

    return {'T':T, 'muQ': muQ, 'muS': muS, 'P':p, 's':s, 'n_B':nB, 'n_Q':nQ, 'n_S':nS, 'n':n, 'e':e, 'chi':chi, 'I':e-3*p}


"""
Helper function to safely convert an array to a new data type.
"""
from __future__ import absolute_import, print_function, division

import numpy as np

import theano

__docformat__ = "restructuredtext en"

def _asarray(a, dtype, order=None):
    """Convert the input to a Numpy array.

    This function is almost identical to ``numpy.asarray``, but it should be
    used instead of its numpy counterpart when a data type is provided in
    order to perform type conversion if required.
    The reason is that ``numpy.asarray`` may not actually update the array's
    data type to the user-provided type. For more information see ticket
    http://projects.scipy.org/numpy/ticket/870.

    In that case, we check that both dtype have the same string
    description (byte order, basic type, and number of bytes), and
    return a view with the desired dtype.

    This function's name starts with a '_' to indicate that it is meant to be
    used internally. It is imported so as to be available directly through
        theano._asarray
    """
    if str(dtype) == 'floatX':
        dtype = theano.config.floatX
    dtype = np.dtype(dtype)  # Convert into dtype object.
    rval = np.asarray(a, dtype=dtype, order=order)
    # Note that dtype comparison must be done by comparing their `num`
    # attribute. One cannot assume that two identical data types are pointers
    # towards the same object (e.g. under Windows this appears not to be the
    # case).
    if rval.dtype.num != dtype.num:
        # Type mismatch between the data type we asked for, and the one
        # returned by numpy.asarray.
        # If both types have the same string description (byte order, basic
        # type, and number of bytes), then it is safe to return a view.
        if (dtype.str == rval.dtype.str):
            # Silent fix.
            return rval.view(dtype=dtype)
        else:
            # Unexpected mismatch: better know what is going on!
            raise TypeError(
                'numpy.array did not return the data type we '
                'asked for (%s %s #%s), instead it returned type '
                '%s %s #%s: function '
                'theano._asarray may need to be modified to handle this '
                'data type.' %
                (dtype, dtype.str, dtype.num, rval.dtype, rval.str,
                 rval.dtype.num))
    else:
        return rval