# -*- coding: utf-8 -*-
from openerp.report import report_sxw
from openerp.report.report_sxw import rml_parse

from datetime import datetime
DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"
DATE_FORMAT = "%Y-%m-%d"

class Parser(rml_parse):

    def __init__(self, cr, uid, name, context):
        super(self.__class__, self).__init__(cr, uid, name, context)
        self.start_date = False
        self.end_date = False
        self.asset_types = []

        self.company_name = False
        self.company_address = False
        self.vat = False
        self.total_residual = False
        self.depreciation_value = False
        self.acc_depreciation = False
        self.remain_value = False

        self.data = {}

        self.localcontext.update({
            'get_vietname_date':self.get_vietname_date,
            'get_header':self.get_header,
            'get_start_date':self.get_start_date,
            'get_end_date':self.get_end_date,
            'get_company_name':self.get_company_name,
            'get_company_address':self.get_company_address,
            'get_company_vat':self.get_company_vat,

            'get_date_type':self.get_date_type,

            'get_wh':self.get_wh,
            'get_dept': self.get_dept,
            'get_line':self.get_line,

            'get_depreciation_value':self.get_depreciation_value,
            'get_acc_depreciation':self.get_acc_depreciation,
            'get_remaining_value':self.get_remaining_value,
        })

    def get_company(self):
        if self.company_id:
            company_obj = self.pool.get('res.company').browse(self.cr, self.uid, self.company_id)
            self.company_name = company_obj.name or ''
            self.company_address = company_obj.street or ''
            self.vat = company_obj.vat or ''
        return True

    def get_company_name(self):
        self.get_header()
        return self.company_name

    def get_company_address(self):
        return self.company_address

    def get_date_type(self):
        if self.date_type == 'depreciation_date':
            return u'Ngày khấu hao'
        if self.date_type == 'using_date':
            return u'Ngày đưa vào sử dụng'
        if self.date_type == 'purchase_date':
            return u'Ngày mua'
        return ''

    def get_company_vat(self):
        return self.vat

    def get_header(self):
        wizard_data = self.localcontext['data']['form']

        self.company_id = wizard_data['company_id'][0]
        self.start_date = wizard_data['date_start']
        self.end_date   = wizard_data['date_end']
        self.date_type = wizard_data['date_type']
        self.show_all = wizard_data['show_all']
        self.asset_types = (','.join(map(str, wizard_data['asset_category_ids'])))
        self.get_company()

        #THANH: Load all data
        res, depreciation_value, acc_depreciation, remain_value = self.pool.get('sql.depreciation.allocation').get_line(self.cr, self.date_type, self.show_all, self.start_date, self.end_date, self.asset_types, self.company_id)
        self.depreciation_value = depreciation_value
        self.acc_depreciation = acc_depreciation
        self.remain_value = remain_value
        self.data = res

    def get_start_date(self):
        return self.get_vietname_date(self.start_date)

    def get_end_date(self):
        return self.get_vietname_date(self.end_date)

    def get_convert_double(self,value):
        return float(value)

    def get_vietname_date(self, date):
        if not date:
            return ''
        date = datetime.strptime(date, DATE_FORMAT)
        return date.strftime('%d/%m/%Y')

    def get_wh(self):
        wh = []
        for key, val in self.data.items():
            wh.append(key)
        return wh

    def get_dept(self, wh):
        dept = []
        for key, val in self.data.items():
            if key == wh:
                for dept_key, dept_vals in val.items():
                    dept.append(dept_key)
        return dept

    def get_line(self, wh, dept):
        return self.data[wh][dept]

    def get_depreciation_value(self):
        return self.depreciation_value

    def get_acc_depreciation(self):
        return self.acc_depreciation

    def get_remaining_value(self):
        return self.remain_value

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import time
import seaborn as sns
from sklearn.feature_selection import f_classif,SelectKBest
from sklearn.model_selection import train_test_split
from imblearn.over_sampling import SMOTE
from sklearn.model_selection import GridSearchCV
from xgboost import XGBClassifier
from sklearn.metrics import confusion_matrix
from sklearn.metrics import recall_score

# Initialising

# This section handles the generic initialisation of maximum rows and columns to be displayed when executed
# This section also sets the path for the datasets

pd.set_option('display.max_rows',None)
pd.set_option('display.max_columns',None)
os.chdir('/Users/praveenthenraj/Desktop/PycharmProjects/Trial/Hackerearth/Novartis/Data/Dataset')

# Reading the data

# This section handles reading the train,test,and sample submission file and creating respective dataframes

train=pd.read_csv('Train.csv')
test=pd.read_csv('Test.csv')
sample=pd.read_csv('sample_submission.csv')
submis=pd.DataFrame()

# Initialising Variables and Data Structures

# This section handles the initialisation of an empty list for storing the important features identified later

imp_feature=[]

# Basic Health Checks

# This section performs the basic health checks on the train and test dataframes:

print('Training Data Shape : '+str(train.shape)+'\n')                                                           #(23856, 18)
print('Testing Data Shape  : '+str(test.shape)+'\n')                                                            #(15903, 17)
print('Training Data Non-int type : '+'\n'+'\n'+str(train.dtypes!=np.int64)+'\n')                               #INCIDENT_ID,DATE,X_12
print('Testing Data Non-int type : '+'\n'+'\n'+str(test.dtypes!=np.int64)+'\n')                                 #INCIDENT_ID,DATE,X_12
print('Training Data Null Columns : '+'\n'+'\n'+str(train.isna().any())+'\n')                                   #X_12
print('Testing Data Null Columns : '+'\n'+'\n'+str(test.isna().any())+'\n')                                     #X_12
print('Count of Null Data in Train : '+'\n'+'\n'+str(train.isna().sum())+'\n')                                  #182
print('Count of Null Data in Test : '+'\n'+'\n'+str(test.isna().sum())+'\n')                                    #127
print('Balance of Dependant Variables : '+'\n'+'\n'+str(train['MULTIPLE_OFFENSE'].value_counts())+'\n')         #0-1068 1-22788
print('No of duplicates in Train Data : '+str(len(train.loc[train['INCIDENT_ID'].duplicated()]==True))+'\n')
print('No of duplicates in Train Data : '+str(len(test.loc[test['INCIDENT_ID'].duplicated()]==True))+'\n')
print('IQR for feature X_12 : '+str(train.describe()))                                                         #0.9740
print('IQR for feature X_12 : '+str(test.describe()))                                                          #0.9722

# EDA

# This section is used to create the EDA for the given datasets

# Plot 1: To identify the distribution of values in X_12 feature to impute for NaN values in train data:
# From the plot, it can be seen that 90% of X_12 value in train_data has value of 1

train['X_12'].value_counts().plot(kind='bar')
plt.xlabel('X_12_Train')
plt.ylabel('Value_counts')
plt.title('Distribution of values in X_12_Train')
plt.show()

# Plot 2: To identify the distribution of each predictor variable and their respective mean and std deviation
# This plot helps to identify the scaling range of each features

train_eda = train.iloc[:, 2:17]

for i in train_eda.columns:
    std = np.round(np.std(train[i]), 2)
    mean = np.round(np.mean(train[i]), 2)
    train[i].value_counts().plot(kind='bar',label=r'Mean    ' + ' - ' + str(mean) + '\n' + r'Std Dev' + ' - ' + str(std))
    plt.xlabel(i)
    plt.ylabel('Value_counts')
    plt.legend()
    plt.title('Distribution of Anonymized Logging parameter ' + i)
    plt.show()

# Plot 3: To identify the correlation between the various dependant variables
# This plots helps to identify the features with multicollinearity. This helps in identifying and dropping those
# features which are highly correlated.

sns.heatmap(train.corr(), annot=True)
plt.show()

# Plot 4: To identify the distribution of values in X_12 feature to impute for NaN values in test data:
# From the plot, it can be seen that the distribution in test data is similar to the train data and most
# value is 1

test['X_12'].value_counts().plot(kind='bar')
plt.xlabel('X_12_Test')
plt.ylabel('Value_counts')
plt.title('Distribution of values in X_12_Test')
plt.show()

# From the above plot it can be seen that the following pairs are highly correlated with each other and hence
# both features can be dropped in case of low variance test score(which will be done later below) or either
# feature will be dropped in case of high variance test score.

# (X_2,X_3)   - correlation of 1 (multicollinear)
# (X_10,X_12) - correlation of 0.88 (multicollinear)

# Feature Preprocessing and Feature Transformation:

# This section handles preprocessing of the input dataset like filling the missing values.
# From the given datasets it is found that only feature X_12 has missing values in both train and test dataset.
# Filling the na values with 1 as Plot_1 suggested that X_12 has maximum number of values as 1 and also a
# mean of values in X_12 in train and test data is 0.9740 and 0.9722 respectively
# We are considering mean to fill Nan values, because from plot1 and plot 4 we can see that the distribution
# of data in feature X_12 is right skewed. Considering the central tendency measure mean for
# such distributions is a better approach.

def feature_preprocess_transformation(data):
    global train,test
    data['X_12']=data['X_12'].fillna(value=1.0).astype(int)

# Calling the preprocessing function for Train and Test Data
# Functions to fill the missing values in train and test dataset

feature_preprocess_transformation(train)
feature_preprocess_transformation(test)

# Feature Importance and Train_Test Data Split
# This section handles identifying the important features and splitting the training data as train_test split
# to train and validate the model

# Here the feature incident_id has been dropped because it does not hold any cause for the cyber attack.
# The feature MULTIPLE_OFFENSE has been dropped as it is the target variable

indep=train.drop(['INCIDENT_ID','DATE','MULTIPLE_OFFENSE'],axis=1).values
dep=train['MULTIPLE_OFFENSE'].values

# Here we are doing an ANNOVA TEST to identify the feature importance with respect to the target variable
# ANNOVA Test is chosen, as the independant variables are continous variables whereas the output variables are
# categorical

x=SelectKBest(score_func=f_classif,k='all').fit(indep,dep)
fs=x.transform(indep)
for i in range(len(x.scores_)):
    print('Feature %d: %f' % (i, x.scores_[i]))
    imp_feature.append(i)

#Result of the ANNOVA test:

#X_1: 4.698824
#X_2 17.782549
#X_3: 17.684310
#X_4: 0.203579
#X_5 0.283907
#X_6 1.489952
#X_7: 1.540444
#X_8: 26.635719
#X_9: 0.013363
#X_10: 1491.658800
#X_11: 266.372662
#X_12: 308.553593
#X_13: 1.011512
#X_14: 5.560878
#X_15: 48.221174

# Selecting only features having annova score of atleast more than 20

print([i for i in range(len(x.scores_)) if x.scores_[i]>20])

# Dropping features having ANNOVA score less than 20 and also features that are correlated based on the EDA done earlier

#'X_1','X_4','X_5','X_6','X_7','X_9','X_13','X_14' --- dropped due to low ANNOVA score
#'X_2','X_12'  --- dropped due to interaction effect between these variables and X_3,X_10 respectively

X=train.drop(['INCIDENT_ID','DATE','X_1','X_2','X_4','X_5','X_6','X_7','X_9','X_12','X_13','X_14','MULTIPLE_OFFENSE'],axis=1).values
y=train['MULTIPLE_OFFENSE'].values

# From the basic health checks done earlier it can be seen that the training dataset is an imbalanced dataset
#Below shows the distribution of each class in the training dataset.
            #0-1068
            #1-22788

# SMOTE technique is selected to upsample the minority target class

sm=SMOTE(sampling_strategy='minority',random_state=2)
X_train_sm,y_train_sm=sm.fit_sample(X,y.ravel())
print(X_train_sm.shape)
print(y_train_sm.shape)
print(y_train_sm.value_counts())

# From above, it can be seen that the minority class '0' has been upsampled

# Splitting the training dataset(after SMOTE) into train-test dataset to train the model and validate the model against
#the test data

X_train,X_test,y_train,y_test=train_test_split(X_train_sm,y_train_sm,test_size=0.2,stratify=y_train_sm,random_state=20)

# Hyperparameter Tuning

# Initially the below set of learning_rate and n_estimatores were tried to find the best values

#l=[0.001,0.005,0.01,0.05,0.1,0.5]
#n=[i for i in range(100,1000,100)]

# From the above set it was identified that the best value for learning_rate is 0.5
# and estimators as 100 for the model.
# Hence learning_rate and n_estimators was narrowed down around 0.5 and 100 respectively and hence
# a GridSearchCV in the range of 0.5-0.8 and 100-120 was applied for learning_rate and n_estimators
# respectively

# Below are the final set of hyperparameters selected to tune the model to identify the best
# possible hyperparameters.

#l=[i for i in np.arange(0.5,0.8,0.01)]
#n=[i for i in range(100,120,1)]
#md=[i for i in np.arange(3,8,1)]
#rl=[i for i in np.arange(0,1,0.1)]

#param={'learning_rate':l}
#param={'n_estimators':n}
#param={'n_estimators':n,'max_depth':md}
#param={'reg_lambda':rl}

# A GridSearchCV was chose for trying out the various combinations of hyperparameters on the model.
# GridSearchCV was selected over other methods like RandomSearch as the intention was not to miss the
# local/global minima by missing a particular parameter during optimisation

# XGBoost model has been selected to classify the datapoints due to the below considerations.
          # Decision tree based models are proven to perform better on classification problems
          # Boosting techniques like XGBoost is preferred over traditional Decision Trees because
          #    in boosting algorithms, we combine a set of weak learners to form a strong model thus
          #    improving the performance of the model.
          # XGBoost model is chosen over other boosting model like Gradient Boosting as XGBoost has
          #    inbuilt L1 and L2 regularization which prevents overfitting of the model which is
          #    prone to happen espicially during classification using decision tree based algorithms.

#xgb1=XGBClassifier()
#xgb=GridSearchCV(xgb1,param_grid=param,cv=5)
#xgb.fit(X_train,y_train)
#print(xgb.best_params_)
#print(xgb.best_estimator_)
#y_pred=xgb.best_estimator_.predict(X_test)

# Model

# Based on the above hyperparameter tuning the below parameters have found to be precise for the given model

xgb=XGBClassifier(learning_rate=0.7200000000000001,max_depth=3,n_estimators=100,reg_lambda=0.1)
xgb.fit(X_train,y_train)
y_pred=xgb.predict(X_test)

tn,fp,fn,tp=confusion_matrix(y_test,y_pred).ravel()
recall=recall_score(y_test,y_pred)
print(recall)
print(tn,fp,fn,tp)

# Processing the test file

X_sub=test.drop(['INCIDENT_ID','DATE','X_1','X_2','X_4','X_5','X_6','X_7','X_9','X_12','X_13','X_14'],axis=1).values                           #Dropping features having Annova score less than 1. X_3 dropped even score is 17 as its score is same as X_2
y_sub=xgb.predict(X_sub)

# Preparing the submission file

submis['INCIDENT_ID']=test['INCIDENT_ID']
submis['MULTIPLE_OFFENSE']=y_sub
submis.to_csv('outfile.csv',index=False)


import tkinter as tk
from tkinter import ttk
import matplotlib
matplotlib.use("TkAgg") #backend
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
from matplotlib.figure import Figure

LARGE_FONT = ("Verdana", 12)

class SpectralGui(tk.Tk):

    def __init__(self,*args,**kwargs):

        tk.Tk.__init__(self, *args, **kwargs)
        tk.Tk.wm_title(self,"Spectral GUI")
        tk.Tk.iconbitmap(self,'c:/users/juanr/downloads/logo_lec.ico')

        container = tk.Frame(self)
        container.pack(side="top", fill="both",expand=True)
        container.grid_rowconfigure(0,weight=1)
        container.grid_columnconfigure(0, weight=1)

        self.frames = {} #DICCIONARIO

        for F in (StartPage,PageOne,PageTwo,PageThree):

            frame = F(container,self)

            self.frames[F] = frame

            frame.grid(row=0, column=0, sticky="nsew")

        self.show_frame(StartPage)

    def show_frame(self,cont):
        frame = self.frames[cont]
        frame.tkraise()

def gf(strintoprint):
    print(strintoprint)

class StartPage(tk.Frame):

    def __init__(self, parent, controller):
        tk.Frame.__init__(self, parent)
        label = tk.Label(self, text = "Start Page", font = LARGE_FONT)
        label.pack(pady=10, padx=10)

        button = ttk.Button(self, text = "Visit Page 1",
                            command = lambda: controller.show_frame(PageOne))
        button.pack()

        button1 = ttk.Button(self, text="Visit Page 2",
                            command=lambda: controller.show_frame(PageTwo))
        button1.pack()

        button2 = ttk.Button(self, text="Visit Graph Page",
                             command=lambda: controller.show_frame(PageThree))
        button2.pack()

class PageOne(tk.Frame):

    def __init__(self,parent,controller):
        tk.Frame.__init__(self, parent)
        label = tk.Label(self, text="Page One", font=LARGE_FONT)
        label.pack(pady=10, padx=10)

        button = ttk.Button(self, text="Back to Home",
                            command=lambda: controller.show_frame(StartPage))
        button.pack()

        button2 = ttk.Button(self, text="Page Two",
                            command=lambda: controller.show_frame(PageTwo))
        button2.pack()

class PageTwo(tk.Frame):

    def __init__(self,parent,controller):
        tk.Frame.__init__(self, parent)
        label = tk.Label(self, text="Page Two", font=LARGE_FONT)
        label.pack(pady=10, padx=10)

        button = ttk.Button(self, text="Back to Home",
                            command=lambda: controller.show_frame(StartPage))
        button.pack()

        button2 = ttk.Button(self, text="Page One",
                           command=lambda: controller.show_frame(PageOne))
        button2.pack()

class PageThree(tk.Frame):

    def __init__(self,parent,controller):
        tk.Frame.__init__(self, parent)
        label = tk.Label(self, text="Graph Page", font=LARGE_FONT)
        label.pack(pady=10, padx=10)

        button = ttk.Button(self, text="Back to Home",
                            command=lambda: controller.show_frame(StartPage))
        button.pack()

        f = Figure(figsize=(5,5),dpi=100)
        a = f.add_subplot(111)
        a.plot([1,2,3,4,5,6,7,8],[1,2,3,4,5,6,7,8])

        #graph
        canvas = FigureCanvasTkAgg(f,self)
        canvas.draw()
        canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand = True)

        #navigation toolbar
        toolbar = NavigationToolbar2Tk(canvas,self)
        toolbar.update()
        #canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=True) #para que este arriba la toolbar
        canvas._tkcanvas.pack(side=tk.BOTTOM, fill=tk.BOTH, expand = True)

app = SpectralGui()
app.mainloop()


from collections import namedtuple
import math
import logging

from .error import EngineError
from . import const
from .flow import coroutine

logger = logging.getLogger(__name__)

Data = namedtuple('Data', 'tissues gf')
Data.__doc__ = """
Data for ZH-L16-GF decompression model.

:var tissues: Tissues gas loading. Tuple of pair numbers - each pair holds
    value of inert gas pressure (N2, He) in a tissue compartment.
:var gf: Gradient factor value.
"""

def eq_gf_limit(gf, p_n2, p_he, a_n2, b_n2, a_he, b_he):
    """
    Calculate ascent ceiling limit of a tissue compartment using Buhlmann
    equation extended with gradient factors by Erik Baker.

    The returned value is absolute pressure of depth of the ascent ceiling.

    :param gf: Gradient factor value.
    :param p_n2: Current tissue pressure for nitrogen.
    :param p_he: Current tissue pressure for helium.
    :param a_n2: Nitrox Buhlmann coefficient A.
    :param b_n2: Nitrox Buhlmann coefficient B.
    :param a_he: Helium Buhlmann coefficient A.
    :param b_he: Helium Buhlmann coefficient B.
    """
    assert gf > 0 and gf <= 1.5
    p = p_n2 + p_he
    a = (a_n2 * p_n2 + a_he * p_he) / p
    b = (b_n2 * p_n2 + b_he * p_he) / p
    return (p - a * gf) / (gf / b + 1 - gf)

class ZH_L16_GF(object):
    """
    Base abstract class for Buhlmann ZH-L16 decompression model with
    gradient factors by Erik Baker - ZH-L16B-GF.

    :var gf_low: Gradient factor low parameter.
    :var gf_high: Gradient factor high parameter.
    :var water_vapour_pressure: Water vapour pressure.
    :var n2_k_const: Gas decay constants :math:`k` for nitrogen for each
        tissue compartment.
    :var he_k_const: Gas decay constants :math:`k` for helium for each
        tissues compartment.
    """
    NUM_COMPARTMENTS = 16
    N2_A = None
    N2_B = None
    HE_A = None
    HE_B = None
    N2_HALF_LIFE = None
    HE_HALF_LIFE = None
    START_P_N2 = 0.7902 # starting pressure of N2 in tissues
    START_P_HE = 0.0    # starting pressure of He in tissues

    def __init__(self):
        """
        Create instance of the model.
        """
        super().__init__()
        self.n2_k_const = self._k_const(self.N2_HALF_LIFE)
        self.he_k_const = self._k_const(self.HE_HALF_LIFE)
        self.gf_low = 0.3
        self.gf_high = 0.85

        self.water_vapour_pressure = const.WATER_VAPOUR_PRESSURE_DEFAULT

    def init(self, surface_pressure):
        """
        Initialize pressure of inert gas in all tissues.

        The method uses starting tissue pressure values for nitrogen and
        helium.

        :param surface_pressure: Surface pressure [bar].
        """
        p_n2 = self.START_P_N2 * (surface_pressure - self.water_vapour_pressure)
        p_he = self.START_P_HE
        data = Data(tuple([(p_n2, p_he)] * self.NUM_COMPARTMENTS), self.gf_low)
        return data

    def load(self, abs_p, time, gas, rate, data):
        """
        Calculate gas loading for all tissue compartments.

        The method returns decompression data model information.

        :param abs_p: Absolute pressure [bar] (current depth).
        :param time: Time of exposure [min] (i.e. time of ascent).
        :param gas: Gas mix configuration.
        :param rate: Pressure rate change [bar/min].
        :param data: Decompression model data.

        .. seealso::

            - :py:meth:`decotengu.model.ZH_L16_GF._tissue_loaders`
            - :py:meth:`decotengu.model.ZH_L16_GF._tissue_loader`
        """
        n2_loader, he_loader = self._tissue_loaders(abs_p, gas, rate)

        tp = tuple(
            (n2_loader(time, p_n2, i), he_loader(time, p_he, i))
            for i, (p_n2, p_he) in enumerate(data.tissues)
        )
        return Data(tp, data.gf)

    def ceiling_limit(self, data, gf=None):
        """
        Calculate pressure of ascent ceiling limit using decompression
        model data.

        The pressure is the shallowest depth a diver can reach without
        decompression sickness. If pressure limit is 3 bar, then diver
        should not go shallower than 20m.

        FIXME: the method signature is gradient factor specific, the
            signature has to be made decompression model independent

        :param data: Decompression model data.
        :param gf: Gradient factor value, `gf_low` by default.

        .. seealso::

            - :py:meth:`decotengu.model.ZH_L16_GF.gf_limit`
            - :py:meth:`decotengu.model.ZH_L16_GF._tissue_loader`
        """
        return max(self.gf_limit(gf, data))

    def _k_const(self, half_life):
        """
        Calculate gas decay constant :math:`k` for each tissue compartment
        half-life value.

        :param half_life: Collection of half-life values for each tissue
            compartment.
        """
        return tuple(const.LOG_2 / v for v in half_life)

    def _exp(self, time, k):
        """
        Calculate value of exponential function for time and gas decay
        constant :math:`k`.

        :param time: Time of exposure [min].
        :param k: Gas decay constant :math:`k` for a tissue compartment.
        """
        return math.exp(-k * time)

    def _tissue_loaders(self, abs_p, gas, rate):
        """
        Create function to load tissue compartment with inert gas for each
        inert gas specified in gas mix configuration.

        :param abs_p: Absolute pressure of current depth [bar] (:math:`P_{abs}`).
        :param gas: Gas mix configuration.
        :param rate: Pressure rate change [bar/min] (:math:`P_{rate}`).
        """
        n2_loader = self._tissue_loader(
            abs_p, gas.n2 / 100, rate, self.n2_k_const
        )
        he_loader = self._tissue_loader(
            abs_p, gas.he / 100, rate, self.he_k_const
        )
        return n2_loader, he_loader

    def _tissue_loader(self, abs_p, f_gas, rate, k_const):
        """
        Create function to load tissue compartment with inert gas.

        The created function uses Schreiner equation and has the following
        parameters

        time
            Time of exposure [min] at depth (:math:`T_{time}`).
        p_i
            Initial (current) pressure of inert gas in tissue compartment
            [bar] (:math:`P_{i}`).
        tissue_no
            Number of tissue compartment in the decompression model
            (starting with zero).

        See :ref:`eq-schreiner` section for details.

        :param abs_p: Absolute pressure of current depth [bar] (:math:`P_{abs}`).
        :param f_gas: Inert gas fraction, i.e. for air it is 0.79 (:math:`F_{gas}`).
        :param rate: Pressure rate change [bar/min] (:math:`P_{rate}`).
        :param k_const: Collection of gas decay constants for each tissue
            compartment (:math:`k`).
        """
        p_alv = f_gas * (abs_p - self.water_vapour_pressure)
        r = f_gas * rate
        def f(time, p_i, tissue_no):
            assert time > 0
            k = k_const[tissue_no]
            return p_alv + r * (time - 1 / k) - (p_alv - p_i - r / k) \
                * self._exp(time, k)
            #return p_alv + r * (t - 1 / k) - (p_alv - p_i - r / k) * math.exp(-k * t)
        return f

    def gf_limit(self, gf, data):
        """
        Calculate pressure of ascent ceiling for each tissue compartment.

        The method returns a tuple of values - a pressure value for each
        tissue compartment.

        :param gf: Gradient factor.
        :param data: Decompression model data.
        """
        # FIXME: make it model independent
        if gf is None:
            gf = self.gf_low
        assert gf > 0 and gf <= 1.5

        data = zip(data.tissues, self.N2_A, self.N2_B, self.HE_A, self.HE_B)
        return tuple(
            eq_gf_limit(gf, p_n2, p_he, n2_a, n2_b, he_a, he_b)
            for (p_n2, p_he), n2_a, n2_b, he_a, he_b in data
        )

class ZH_L16B_GF(ZH_L16_GF): # source: gfdeco.f by Baker
    """
    ZH-L16B-GF decompression model.
    """
    N2_A = (
        1.1696, 1.0000, 0.8618, 0.7562, 0.6667, 0.5600, 0.4947, 0.4500,
        0.4187, 0.3798, 0.3497, 0.3223, 0.2850, 0.2737, 0.2523, 0.2327,
    )
    N2_B = (
        0.5578, 0.6514, 0.7222, 0.7825, 0.8126, 0.8434, 0.8693, 0.8910,
        0.9092, 0.9222, 0.9319, 0.9403, 0.9477, 0.9544, 0.9602, 0.9653,
    )
    HE_A = (
        1.6189, 1.3830, 1.1919, 1.0458, 0.9220, 0.8205, 0.7305, 0.6502,
        0.5950, 0.5545, 0.5333, 0.5189, 0.5181, 0.5176, 0.5172, 0.5119,
    )
    HE_B = (
        0.4770, 0.5747, 0.6527, 0.7223, 0.7582, 0.7957, 0.8279, 0.8553,
        0.8757, 0.8903, 0.8997, 0.9073, 0.9122, 0.9171, 0.9217, 0.9267,
    )
    N2_HALF_LIFE = (
        5.0, 8.0, 12.5, 18.5, 27.0, 38.3, 54.3, 77.0, 109.0,
        146.0, 187.0, 239.0, 305.0, 390.0, 498.0, 635.0,
    )
    HE_HALF_LIFE = (
        1.88, 3.02, 4.72, 6.99, 10.21, 14.48, 20.53, 29.11,
        41.20, 55.19, 70.69, 90.34, 115.29, 147.42, 188.24, 240.03,
    )

class ZH_L16C_GF(ZH_L16_GF): # source: ostc firmware code
    """
    ZH-L16C-GF decompression model.
    """
    N2_A = (
        1.2599, 1.0000, 0.8618, 0.7562, 0.6200, 0.5043, 0.4410, 0.4000,
        0.3750, 0.3500, 0.3295, 0.3065, 0.2835, 0.2610, 0.2480, 0.2327,
    )
    N2_B = (
        0.5050, 0.6514, 0.7222, 0.7825, 0.8126, 0.8434, 0.8693, 0.8910,
        0.9092, 0.9222, 0.9319, 0.9403, 0.9477, 0.9544, 0.9602, 0.9653,
    )
    HE_A = (
        1.7424, 1.3830, 1.1919, 1.0458, 0.9220, 0.8205, 0.7305, 0.6502,
        0.5950, 0.5545, 0.5333, 0.5189, 0.5181, 0.5176, 0.5172, 0.5119,
    )
    HE_B = (
        0.4245, 0.5747, 0.6527, 0.7223, 0.7582, 0.7957, 0.8279, 0.8553,
        0.8757, 0.8903, 0.8997, 0.9073, 0.9122, 0.9171, 0.9217, 0.9267,
    )
    N2_HALF_LIFE = (
        4.0, 8.0, 12.5, 18.5, 27.0, 38.3, 54.3, 77.0, 109.0,
        146.0, 187.0, 239.0, 305.0, 390.0, 498.0, 635.0,
    )
    HE_HALF_LIFE = (
        1.51, 3.02, 4.72, 6.99, 10.21, 14.48, 20.53, 29.11, 41.20,
        55.19, 70.69, 90.34, 115.29, 147.42, 188.24, 240.03,
    )

class DecoModelValidator(object):
    """
    Dive step tissue pressure validator (coroutine class).

    Create coroutine object, then call it to start the coroutine.

    :var engine: DecoTengu decompression engine.
    """
    def __init__(self, engine):
        """
        Create coroutine object.

        :param engine: DecoTengu decompression engine.
        """
        self.engine = engine
        self._first_stop_checked = False

    @coroutine
    def __call__(self):
        """
        Start the coroutine.
        """
        logger.info('started deco model validator')
        prev = None
        while True:
            step = yield
            self._ceiling_limit(step)
            self._first_stop_at_ceiling(prev, step)
            prev = step

    def _ceiling_limit(self, step):
        """
        Verify that a dive step is deeper than a pressure ceiling limit.

        :param step: Dive step to verify.
        """
        limit = self.engine.model.ceiling_limit(step.data, step.data.gf)
        if step.abs_p < limit: # ok when step.abs_p >= limit
            raise EngineError(
                'Pressure ceiling validation error at {} (limit={})'
                .format(step, limit)
            )

    def _first_stop_at_ceiling(self, prev, step):
        """
        Verify that first decompression stop is at pressure ceiling limit.

        :param prev: Previous dive step.
        :param step: Dive step to verify.
        """
        # FIXME: Phase circular import, so using 'deco_stop' below
        if not self._first_stop_checked and step.phase == 'deco_stop':
            stop = prev
            limit = self.engine.model.ceiling_limit(stop.data)
            # if further ascent was possible, then first deco stop is at
            # wrong depth, i.e. stop at 21m and limit at 17.9 results in
            # error
            if stop.abs_p - self.engine._p3m >= limit:
                raise EngineError(
                    'First decompression stop not at deco ceiling. Error for'
                    ' {} (next step possible to {}, its limit is {})'
                    .format(stop, stop.abs_p - self.engine._p3m, limit)
                )
            self._first_stop_checked = True
            logger.debug('first deco stop ok')

# vim: sw=4:et:ai


#!/usr/bin/python3
#-*- coding: utf-8 -*-

"""import os

from reportlab.lib import colors
from reportlab.lib.pagesizes import letter, inch
from reportlab.platypus import SimpleDocTemplate, Table, TableStyle

doc_save_path = os.getcwd()

if not os.path.exists(doc_save_path+"/prueba"):
    os.mkdir(doc_save_path+"/prueba")

doc = SimpleDocTemplate(doc_save_path+"/prueba/simple_table_grid3.pdf", pagesize=letter)
# container for the 'Flowable' objects
elements = []

data= [['00', '01', '02', '03', '04'],
       ['10', '11', '12', '13', '14'],
       ['20', '21', '22', '23', '24'],
       ['30', '31', '32', '33', '34']]

t=Table(data,5*[0.4*inch], 4*[0.4*inch])

t.setStyle(TableStyle([('ALIGN',(1,1),(-2,-2),'RIGHT'),
                       ('TEXTCOLOR',(1,1),(-2,-2),colors.red),
                       ('VALIGN',(0,0),(0,-1),'TOP'),
                       ('TEXTCOLOR',(0,0),(0,-1),colors.blue),
                       ('ALIGN',(0,-1),(-1,-1),'CENTER'),

                       ('VALIGN',(0,-1),(-1,-1),'MIDDLE'),
                       ('TEXTCOLOR',(0,-1),(-1,-1),colors.green),
                       ('INNERGRID', (0,0), (-1,-1), 0.25, colors.black),
                       ('BOX', (0,0), (-1,-1), 0.25, colors.black),
                       ]))

elements.append(t)
# write the document to disk
doc.build(elements)"""

from reportlab.lib import colors
from reportlab.lib.pagesizes import letter, inch, landscape
from reportlab.pdfgen import textobject
from reportlab.platypus import Image, Paragraph, SimpleDocTemplate, Table
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.pdfgen import canvas
import os

"""P0 = Paragraph('''
               A paragraph
               <super>1</super>''',
               styleSheet["BodyText"])

P = Paragraph('''
    <para align=center spaceb=3>The ReportLab Left
    Logo
    Image</para>''',
    styleSheet["BodyText"])"""

"""data= [['A', 'B', 'C', P0, 'D'],
       ['00', '01', '02', [P], '04'],
       ['10', '11', '12', [P], '14'],
       ['20', '21', '22', '23', '24'],
       ['30', '31', '32', '33', '34']]

t=Table(data,style=[('GRID',(1,1),(-2,-2),1,colors.green),
                    ('BOX',(0,0),(1,-1),2,colors.red),
                    ('LINEABOVE',(1,2),(-2,2),1,colors.blue),
                    ('LINEBEFORE',(2,1),(2,-2),1,colors.pink),
                    ('BACKGROUND', (0, 0), (0, 1), colors.pink),
                    ('BACKGROUND', (1, 1), (1, 2), colors.lavender),
                    ('BACKGROUND', (2, 2), (2, 3), colors.orange),
                    ('BOX',(0,0),(-1,-1),2,colors.black),
                    ('GRID',(0,0),(-1,-1),0.5,colors.black),
                    ('VALIGN',(3,0),(3,0),'BOTTOM'),
                    ('BACKGROUND',(3,0),(3,0),colors.limegreen),
                    ('BACKGROUND',(3,1),(3,1),colors.khaki),
                    ('ALIGN',(3,1),(3,1),'CENTER'),
                    ('BACKGROUND',(3,2),(3,2),colors.beige),
                    ('ALIGN',(3,2),(3,2),'LEFT'),
])
t._argW[3]=1.5*inch

elements.append(t)
# write the document to disk
doc.build(elements)
"""
class saveReport():

    ecuacion = None
    intervalos = None
    tolerancia = None
    raiz = None
    data = None
    metodo = None
    iteraciones = None
    path_biseccion = os.getcwd()+"/Soluciones/Biseccion"
    path_punto_fijo = os.getcwd()+"/Soluciones/Punto_Fijo"

    def __init__(self, eq, inter, tol, raiz_, data_, met, itera):
        self.ecuacion = eq
        self.intervalos = inter
        self.tolerancia = tol
        self.raiz = raiz_
        self.data = data_
        self.metodo = met
        self.iteraciones = itera

        self.crear_carpeta()
        self.crear_carpeta()

    def crear_pdf_biseccion(self):

        header = None

        titulo = "Int_"+str(self.intervalos)+"_Tol_"+self.tolerancia

        if self.metodo == "Biseccion":
            header = "f(x) = "+ self.ecuacion
            titulo = "/Soluciones/Biseccion/"+titulo
        else:
            header = "g(x) = "+ self.ecuacion
            titulo = "/Soluciones/Punto_Fijo/"+titulo

        doc = SimpleDocTemplate(os.getcwd()+titulo+".pdf", pagesize=landscape(letter))
        elements = []

        styleSheet = getSampleStyleSheet()

        header_ = Paragraph("Método de "+self.metodo+"", styleSheet['Heading1'])
        texto1 = Paragraph("La Ecuación "+header+"", styleSheet['Heading2'])
        texto2 = Paragraph("Tiene una Raiz x = "+self.raiz+"", styleSheet['Heading2'])
        texto3 = Paragraph("En el intervalo "+str(self.intervalos)+"", styleSheet['Heading2'])
        texto4 = Paragraph("Con tolerancia "+str(self.tolerancia)+"", styleSheet['Heading2'])
        texto5 = Paragraph("Tardó "+str(self.iteraciones)+" iteraciones en encontrar dicha raíz.", styleSheet['Heading2'])

        elements.append(header_)
        elements.append(texto1)
        elements.append(texto2)
        elements.append(texto3)
        elements.append(texto4)
        elements.append(texto5)
        elements.append(Paragraph("",styleSheet['Heading1']))

        self.data.insert(0, (Paragraph("n", styleSheet["Heading2"]), Paragraph("a",styleSheet["Heading2"]),
                            Paragraph("b",styleSheet["Heading2"]), Paragraph("p",styleSheet["Heading2"]),
                            Paragraph("f(p)",styleSheet["Heading2"])))

        for i in range(1,len(self.data),1):
            self.data[i].insert(0,i)

        t=Table(self.data,style=[
                    ('GRID',(0,0),(-1,-1),1,colors.black),
                    ('BACKGROUND', (0, 0), (0, -1), colors.blue),
                    ('TEXTCOLOR',(0,0),(-1,0),colors.white),
                    ('TEXTCOLOR',(0,0),(0,-1),colors.white),
                    ('SIZE', (0,0), (-1,-1), 15),
                    ('VALIGN',(0,0),(-1,-1),'TOP'),
                    ('BACKGROUND', (0, 0), (-1, 0), colors.red),
                    ('SIZE', (1,1),(-1,-1), 12)
        ])

        t._argW[0]=0.6*inch
        t._argW[1]=1.9*inch
        t._argW[2]=1.9*inch
        t._argW[3]=1.9*inch
        t._argW[4]=2.5*inch

        for i in range(len(self.data)):
            t._argH[i]=0.4*inch

        elements.append(t)
        # write the document to disk
        doc.build(elements)

    def crear_carpeta(self):

        if not os.path.exists(os.getcwd()+"/Soluciones"):
            os.mkdir(os.getcwd()+"/Soluciones")
            os.mkdir(self.path_biseccion)
            os.mkdir(self.path_punto_fijo)

    def crear_pdf_punto(self):

        header = None

        titulo = "Xi_"+str(self.intervalos)+"_Tol_"+self.tolerancia

        if self.metodo == "Biseccion":
            header = "f(x) = "+ self.ecuacion
            titulo = "/Soluciones/Biseccion/"+titulo
        else:
            header = "g(x) = "+ self.ecuacion
            titulo = "/Soluciones/Punto_Fijo/"+titulo

        doc = SimpleDocTemplate(os.getcwd()+titulo+".pdf", pagesize=landscape(letter))
        elements = []

        styleSheet = getSampleStyleSheet()

        header_ = Paragraph("Método de "+self.metodo+"", styleSheet['Heading1'])
        texto1 = Paragraph("La Ecuación "+header+"", styleSheet['Heading2'])
        texto2 = Paragraph("Tiene una Raiz x = "+self.raiz+"", styleSheet['Heading2'])
        texto3 = Paragraph("Valor Inicial   X1 ="+str(self.intervalos)+"", styleSheet['Heading2'])
        texto4 = Paragraph("Con tolerancia "+str(self.tolerancia)+"", styleSheet['Heading2'])
        texto5 = Paragraph("Tardó "+str(self.iteraciones)+" iteraciones en encontrar dicha raíz.", styleSheet['Heading2'])

        elements.append(header_)
        elements.append(texto1)
        elements.append(texto2)
        elements.append(texto3)
        elements.append(texto4)
        elements.append(texto5)
        elements.append(Paragraph("",styleSheet['Heading1']))

        self.data.insert(0, (Paragraph("n", styleSheet["Heading2"]), Paragraph("Xn",styleSheet["Heading2"]),
                            Paragraph("X",styleSheet["Heading2"])))

        for i in range(1,len(self.data),1):
            self.data[i].insert(0,i)

        t=Table(self.data,style=[
                    ('GRID',(0,0),(-1,-1),1,colors.black),
                    ('BACKGROUND', (0, 0), (0, -1), colors.blue),
                    ('TEXTCOLOR',(0,0),(-1,0),colors.white),
                    ('TEXTCOLOR',(0,0),(0,-1),colors.white),
                    ('SIZE', (0,0), (-1,-1), 15),
                    ('VALIGN',(0,0),(-1,-1),'TOP'),
                    ('BACKGROUND', (0, 0), (-1, 0), colors.red),
                    ('SIZE', (1,1),(-1,-1), 12)
        ])

        t._argW[0]=0.6*inch
        t._argW[1]=2.5*inch
        t._argW[2]=2.5*inch

        for i in range(len(self.data)):
            t._argH[i]=0.4*inch

        elements.append(t)
        # write the document to disk
        doc.build(elements)


import pickle
from omegaconf import OmegaConf
from skilltranslation.envs.blockstacking import BlockStackFloatPandaEnv
import gym
import numpy as np
from mani_skill2.utils.wrappers import ManiSkillActionWrapper, NormalizeActionWrapper
import os.path as osp
from skilltranslation.envs.blockstacking.env import BlockStackMagicPandaEnv
from skilltranslation.utils.sampling import resample_teacher_trajectory

def animate(imgs, filename="animation.mp4", _return=True, fps=10):
    if isinstance(imgs, dict):
        imgs = imgs["image"]
    print(f"animating {filename}")
    from moviepy.editor import ImageSequenceClip

    imgs = ImageSequenceClip(imgs, fps=fps)
    imgs.write_videofile(filename, fps=fps)
    if _return:
        from IPython.display import Video

        return Video(filename, embed=True)

def collect_magic_teacher_traj(env: BlockStackMagicPandaEnv, render=False, mode='human', change_in_height=0, change_in_goal_height=0):
    assert env.connected == False

    assert env.goal_coords is not None
    init_xyz = env.get_articulations()[0].get_qpos()[:3]
    magic_traj=dict(
        observations=[],
        attns=[],
    )
    vids=[]
    if render and mode == 'human':
        env.render()
        env.get_viewer().paused=True
    i_step=0

    def goto(xyz, gripper_action, attn=0):
        assert gripper_action is None
        nonlocal i_step, block
        target=np.array(xyz)
        observations=[]  # during this goto()
        attns = []
        while True:
            robot_qpos=env.get_articulations()[0].get_qpos()
            delta_xyz=target - robot_qpos[:3]
            robot_to_target_dist=np.linalg.norm(delta_xyz)
            vel=np.linalg.norm(env.unwrapped._agent._robot.get_qvel()[:3])
            if robot_to_target_dist < 0.01 and vel < 0.05:
                break
            action=np.zeros(4)
            action[:3]=delta_xyz.copy()
            # action[-1]=gripper_action
            # print(action)
            if 'NormalizeActionWrapper' in str(env):
                action[:3]=action[:3] / env.env.action_space.high[0]
                # gripper always +1
                action[-1]=+1
            # collect observations
            obs_dict=env.get_obs()
            dense_obs=np.zeros(12)
            dense_obs[:5]=obs_dict['agent']['qpos']
            # dense_obs[5]=obs_dict['agent']['qvel']
            dense_obs[5:]=obs_dict['extra']['obs_blocks']['absolute'][attn * 7: attn * 7 + 7]
            observations.append(dense_obs)
            attns.append(attn)
            # attentions.append(attn)
            # print("{}: dist {:.4f}, act {}, cur xyz {}, target xyz {}".format(
            #     i_step, robot_to_target_dist, action, robot_qpos[:3], xyz
            # ))
            assert action[-1] == +1
            env.step(action)

            if render:
                if mode == 'rgb_array':
                    vids.append(env.render(mode))
                else:
                    env.render(mode)
            i_step+=1
        magic_traj['observations']+=observations
        magic_traj["attns"] += attns

    for i in range(len(env.blocks)):
        block=env.blocks[i]
        block_pos=env.blocks[i].pose.p
        # above_pos=[block_pos[0], block_pos[1], 0.25]
        # grasp_pos=[block_pos[0], block_pos[1], 0.04 + 0.112]
        above_pos=[block_pos[0], block_pos[1], 0.25+change_in_height]
        grasp_pos=[block_pos[0], block_pos[1], 0.113]

        goto(xyz=above_pos, gripper_action=None, attn=i)
        goto(xyz=grasp_pos, gripper_action=None, attn=i)
        env.magic_grasp(block)
        goto(xyz=above_pos, gripper_action=None, attn=i)

        goal_pos=env.goal_coords[i]
        # goal_above_pos=[goal_pos[0], goal_pos[1], 0.25]
        # release_pos=[goal_pos[0], goal_pos[1], goal_pos[2] + 0.132]
        goal_above_pos=[goal_pos[0], goal_pos[1], goal_pos[2]+0.25+ change_in_goal_height]
        release_pos=[goal_pos[0], goal_pos[1], goal_pos[2] + 0.096]

        goto(xyz=goal_above_pos, gripper_action=None, attn=i)
        goto(xyz=release_pos, gripper_action=None, attn=i)
        env.magic_release()
        goto(xyz=goal_above_pos, gripper_action=None, attn=i)

        # back to init
        print("go back to init")
        goto(xyz = init_xyz, gripper_action=None, attn=i)
        print("finished init")
        # last observation and attention:
    obs_dict=env.get_obs()
    dense_obs=np.zeros(12)
    dense_obs[:5]=obs_dict['agent']['qpos']
    dense_obs[5:]=obs_dict['extra']['obs_blocks']['absolute'][i * 7:i * 7 + 7]
    magic_traj['observations'].append(dense_obs)

    magic_traj['attns'] = np.array(magic_traj['attns'], dtype=int)
    magic_traj['observations'] = np.array(magic_traj['observations'], dtype=np.float32)
    # assert len(float_traj['attentions']) == len(float_traj['observations'])
    if render and mode == 'rgb_array':
        return magic_traj, vids

    if not render:
        return magic_traj
    return magic_traj

def collect_teacher_traj(env: BlockStackFloatPandaEnv, render=False, mode='human', fixed_gripper=True):
    assert env.goal_coords is not None
    float_traj = dict(
        observations=[],
        attns=[],
    )
    vids = []
    init_xyz = env.get_articulations()[0].get_qpos()[:3]
    if render and mode=='human':
        env.render()
        env.get_viewer().paused = True

    i_step = 0

    def goto(xyz, gripper_action, attn=0):
        nonlocal i_step, block
        target = np.array(xyz)
        observations = []
        attns = []
        while True:
            ee_pos = env.unwrapped._agent.hand_link.get_pose().p
            delta_xyz = target - ee_pos
            robot_to_target_dist = np.linalg.norm(delta_xyz)
            vel = np.linalg.norm(env.unwrapped._agent._robot.get_qvel()[:3])

            grasp_flag = env.unwrapped._agent.check_grasp(block)
            gripper_flag = grasp_flag != (gripper_action > 0)

            if robot_to_target_dist < 0.01 and vel<0.05 and gripper_flag:
                break
            action = np.zeros(4)
            action[:3] = delta_xyz.copy()
            action[-1] = gripper_action
            # print(action)
            if 'NormalizeActionWrapper' in str(env):
                action[:3] = action[:3] / env.env.action_space.high[0]
                action[-1] = 1 if gripper_action > 0 else -1
            obs_dict=env.get_obs()
            dense_obs=np.zeros(12)
            dense_obs[:5]=obs_dict['agent']['qpos']
            if fixed_gripper:
                dense_obs[3:5] = np.array([0.04, 0.04])
            # dense_obs[5:10]=obs_dict['agent']['qvel']
            dense_obs[5:]=obs_dict['extra']['obs_blocks']['absolute'][attn * 7: attn * 7 + 7]
            observations.append(dense_obs)
            attns.append(attn)
            # print("{}: dist {:.4f}, act {}, cur xyz {}, target xyz {}".format(
            #     i_step, robot_to_target_dist, action, ee_pos, xyz
            # ))
            # print(env.unwrapped._agent.hand_link.get_pose())
            env.step(action)

            if render:
                if mode == 'rgb_array':
                    vids.append(env.render(mode))
                else:
                    env.render(mode)
            i_step += 1
        float_traj["observations"] += observations
        float_traj["attns"] += attns

    for i in range(len(env.blocks)):
        block = env.blocks[i]
        block_pos = env.blocks[i].pose.p
        goal_pos = env.goal_coords[i]
        transport_height = goal_pos[2] + 0.12

        above_pos = [block_pos[0], block_pos[1], 0.25]
        grasp_pos = [block_pos[0], block_pos[1], 0.113]
        start_transport_pos = [block_pos[0], block_pos[1], transport_height]
        end_transport_pos = [goal_pos[0] - 0.04, goal_pos[1], transport_height]

        goto(xyz=above_pos, gripper_action=0.04, attn=i)
        goto(xyz=grasp_pos, gripper_action=0.04, attn=i)
        goto(xyz=grasp_pos, gripper_action=0, attn=i)
        goto(xyz=start_transport_pos, gripper_action=0, attn=i)
        goto(xyz=end_transport_pos, gripper_action=0, attn=i)

        goal_above_pos = [goal_pos[0], goal_pos[1], transport_height]
        release_pos = [goal_pos[0], goal_pos[1], goal_pos[2] + 0.096]

        goto(xyz=goal_above_pos, gripper_action=0, attn=i)
        goto(xyz=release_pos, gripper_action=0, attn=i)
        goto(xyz=release_pos, gripper_action=0.04, attn=i)
        goto(xyz=goal_above_pos, gripper_action=0.04, attn=i)
        goto(xyz=end_transport_pos, gripper_action=0.04, attn=i)
        goto(xyz = init_xyz, gripper_action=0.04, attn=i)

    # if render and mode == 'rgb_array':
    #     return vids
    float_traj['attns'] = np.array(float_traj['attns'], dtype=int)
    float_traj['observations'] = np.array(float_traj['observations'], dtype=np.float32)
    return float_traj

if __name__ == '__main__':
    cfg = OmegaConf.from_cli()
    from tqdm import tqdm
    np.set_printoptions(suppress=True, precision=3)
    save_path = cfg.save_path
    # size = cfg.size
    goal = cfg.goal
    N = 2
    if "n" in cfg:
        N = cfg.n
    env_name = 'BlockStackMagic-v0'
    env = gym.make(
        env_name,
        reward_mode='sparse',
        obs_mode='state_dict',
        num_blocks=1 if "train" in goal else -1,
        goal=goal,
    )
    env = ManiSkillActionWrapper(env)
    env = NormalizeActionWrapper(env)
    dataset = dict(teacher=dict())
    for i in tqdm(range(N)):
        env.reset(seed=i)
        # change_height = np.random.uniform(0,0.36)
        # change_goal_height = np.random.uniform(0,0.15)
        traj = collect_magic_teacher_traj(env, render=True, mode='human', change_in_goal_height=0.1, change_in_height=0.18)
        max_block_id = traj["attns"].max()
        all_obs = []
        attns = []
        for b_id in sorted(np.unique(traj["attns"])):
            obs = traj["observations"][np.where(traj["attns"] == b_id)]
            obs = resample_teacher_trajectory(obs, max_dist=4e-2)
            dataset["teacher"][f"{i}-{b_id}"] = dict(observations=obs)
            attns.append(np.ones(len(obs), dtype=int)*b_id)
            all_obs.append(obs)
        all_obs = np.concatenate(all_obs)
        attns = np.concatenate(attns)
        dataset["teacher"][f"{i}"] = dict(observations=all_obs, attns=attns)
        # do even resampling of teacher trajectory.
        # traj = resample_teacher_trajectory(traj["observations"])
        # dataset["traj"]
        # import pdb;pdb.set_trace()
        # vids = collect_teacher_traj(env, render=True, mode='rgb_array')
        # animate(vids, 'videos/teacher_{:d}.mp4'.format(i), fps=24)
    with open(osp.join(save_path), "wb") as f:
        pickle.dump(dataset, f)


#!/usr/bin/env python
"""

Neural SPARQL Machines - Generator utils.

'SPARQL as a Foreign Language' by Tommaso Soru and Edgard Marx et al., SEMANTiCS 2017
https://arxiv.org/abs/1708.07624

Version 2.0.0

"""
import collections
import http.client
import json
import logging
import re
import sys
import urllib.request, urllib.parse, urllib.error
import urllib.request, urllib.error, urllib.parse
from functools import reduce

ENDPOINT = "http://dbpedia.org/sparql"
GRAPH = "http://dbpedia.org"

def log_statistics( used_resources, special_classes, not_instanced_templates ):
    total_number_of_resources = len(used_resources)
    total_number_of_filled_placeholder_positions = sum(used_resources.values())
    examples_per_instance = collections.Counter()
    for resource in used_resources:
        count = used_resources[resource]
        examples_per_instance.update([count])

    logging.info('{:6d} used resources in {} placeholder positions'.format(total_number_of_resources, total_number_of_filled_placeholder_positions))
    for usage in examples_per_instance:
        logging.info('{:6d} resources occur \t{:6d} times \t({:6.2f} %) '.format(examples_per_instance[usage], usage, examples_per_instance[usage]*100/total_number_of_resources))
    for cl in special_classes:
        logging.info('{} contains: {}'.format(cl, ', '.join(special_classes[cl])))
    logging.info('{:6d} not instanciated templates:'.format(sum(not_instanced_templates.values())))
    for template in not_instanced_templates:
        logging.info('{}'.format(template))

def save_cache ( file, cache ):
    ordered = collections.OrderedDict(cache.most_common())
    with open(file, 'w') as outfile:
        json.dump(ordered, outfile)

def query_dbpedia( query ):
    param = dict()
    param["default-graph-uri"] = GRAPH
    param["query"] = query
    param["format"] = "JSON"
    param["CXML_redir_for_subjs"] = "121"
    param["CXML_redir_for_hrefs"] = ""
    param["timeout"] = "600"
    param["debug"] = "on"
    try:
        resp = urllib.request.urlopen(ENDPOINT + "?" + urllib.parse.urlencode(param))
        j = resp.read()
        resp.close()
    except (urllib.error.HTTPError, http.client.BadStatusLine):
        logging.debug("*** Query error. Empty result set. ***")
        j = '{ "results": { "bindings": [] } }'
    sys.stdout.flush()
    return json.loads(j)

def strip_brackets(s):
    s = re.sub(r'\([^)]*\)', '', s)
    if "," in s:
        s = s[:s.index(",")]
    return s.strip().lower()

REPLACEMENTS = [
    ['dbo:', 'http://dbpedia.org/ontology/', 'dbo_'],
    ['dbp:', 'http://dbpedia.org/property/', 'dbp_'],
    ['dbc:', 'http://dbpedia.org/resource/Category:', 'dbc_'],
    ['dbr:', 'res:', 'http://dbpedia.org/resource/', 'dbr_'],
    ['dct:', 'dct_'],
    ['geo:', 'geo_'],
    ['georss:', 'georss_'],
    ['rdf:', 'rdf_'],
    ['rdfs:', 'rdfs_'],
    ['foaf:', 'foaf_'],
    ['owl:', 'owl_'],
    ['yago:', 'yago_'],
    ['skos:', 'skos_'],
    [' ( ', '  par_open  '],
    [' ) ', '  par_close  '],
    ['(', ' attr_open '],
    [') ', ')', ' attr_close '],
    ['{', ' brack_open '],
    ['}', ' brack_close '],
    [' . ', ' sep_dot '],
    ['. ', ' sep_dot '],
    ['?', 'var_'],
    ['*', 'wildcard'],
    [' <= ', ' math_leq '],
    [' >= ', ' math_geq '],
    [' < ', ' math_lt '],
    [' > ', ' math_gt ']
]

STANDARDS = {
        'dbo_almaMater': ['dbp_almaMater'],
        'dbo_award': ['dbp_awards'],
        'dbo_birthPlace': ['dbp_birthPlace', 'dbp_placeOfBirth'],
        'dbo_deathPlace': ['dbp_deathPlace', 'dbp_placeOfDeath'],
        'dbo_child': ['dbp_children'],
        'dbo_college': ['dbp_college'],
        'dbo_hometown': ['dbp_hometown'],
        'dbo_nationality': ['dbo_stateOfOrigin'],
        'dbo_relative': ['dbp_relatives'],
        'dbo_restingPlace': ['dbp_restingPlaces', 'dbp_placeOfBurial', 'dbo_placeOfBurial', 'dbp_restingplace'],
        'dbo_spouse': ['dbp_spouse'],
        'dbo_partner': ['dbp_partner']
}

def encode( sparql ):
    encoded_sparql = do_replacements(sparql)
    shorter_encoded_sparql = shorten_query(encoded_sparql)
    normalized = normalize_predicates(shorter_encoded_sparql)
    return normalized

def decode ( encoded_sparql ):
    short_sparql = reverse_replacements(encoded_sparql)
    sparql = reverse_shorten_query(short_sparql)
    return sparql

def normalize_predicates( sparql ):
    for standard in STANDARDS:
        for alternative in STANDARDS[standard]:
            sparql = sparql.replace(alternative, standard)

    return sparql

def do_replacements( sparql ):
    for r in REPLACEMENTS:
        encoding = r[-1]
        for original in r[:-1]:
            sparql = sparql.replace(original, encoding)
    return sparql

def reverse_replacements( sparql ):
    for r in REPLACEMENTS:
        original = r[0]
        encoding = r[-1]
        sparql = sparql.replace(encoding, original)
        stripped_encoding = str.strip(encoding)
        sparql = sparql.replace(stripped_encoding, original)
    return sparql

def shorten_query( sparql ):
    sparql = re.sub(r'order by desc\s+....?_open\s+([\S]+)\s+....?_close', '_obd_ \\1', sparql, flags=re.IGNORECASE)
    sparql = re.sub(r'order by asc\s+....?_open\s+([\S]+)\s+....?_close', '_oba_ \\1', sparql, flags=re.IGNORECASE)
    sparql = re.sub(r'order by\s+([\S]+)', '_oba_ \\1', sparql, flags=re.IGNORECASE)
    return sparql

def reverse_shorten_query( sparql ):
    sparql = re.sub(r'_oba_ ([\S]+)', 'order by asc (\\1)', sparql, flags=re.IGNORECASE)
    sparql = re.sub(r'_obd_ ([\S]+)', 'order by desc (\\1)', sparql, flags=re.IGNORECASE)
    return sparql

def read_template_file(file):
    annotations = list()
    line_number = 1
    with open(file) as f:
        for line in f:
            values = line[:-1].split(';')
            target_classes = [values[0] or None, values[1] or None, values[2] or None]
            question = values[3]
            query = values[4]
            generator_query = values[5]
            id = values[6] if (len(values) >= 7 and values[6]) else line_number
            line_number += 1
            annotation = Annotation(question, query, generator_query, id, target_classes)
            annotations.append(annotation)
    return annotations

class Annotation:
    def __init__(self, question, query, generator_query, id=None, target_classes=None):
        self.question = question
        self.query = query
        self.generator_query = generator_query
        self.id = id
        self.target_classes = target_classes if target_classes != None else []
        self.variables = extract_variables(generator_query)

def extract_variables(query):
    variables = []
    query_form_pattern = r'^.*?where'
    query_form_match = re.search(query_form_pattern, query, re.IGNORECASE)
    if query_form_match:
        letter_pattern = r'\?(\w)'
        variables = re.findall(letter_pattern, query_form_match.group(0))
    return variables

def extract_encoded_entities( encoded_sparql ):
    sparql = decode(encoded_sparql)
    entities = extract_entities(sparql)
    encoded_entities = list(map(encode, entities))
    return encoded_entities

def extract_entities( sparql ):
    triples = extractTriples(sparql)
    entities = set()
    for triple in triples:
        possible_entities = [triple['subject'], triple['object']]
        sorted_out = [e for e in possible_entities if not e.startswith('?') and ':' in e]
        entities = entities.union([re.sub(r'^optional{', '', e, flags=re.IGNORECASE) for e in sorted_out])
    return entities

def extract_predicates( sparql ):
    triples = extractTriples(sparql)
    predicates = set()
    for triple in triples:
        pred = triple['predicate']
        predicates.add(pred)
    return predicates

def extractTriples (sparqlQuery):
    triples = []
    whereStatementPattern = r'where\s*?{(.*?)}'
    whereStatementMatch = re.search(whereStatementPattern, sparqlQuery, re.IGNORECASE)
    if whereStatementMatch:
        whereStatement = whereStatementMatch.group(1)
        triples = splitIntoTriples(whereStatement)
    return triples

def splitIntoTriples (whereStatement):
    tripleAndSeparators = re.split(r'(\.[\s\?\<$])', whereStatement)
    trimmed = [str.strip() for str in tripleAndSeparators]

    def repair (list, element):
        if element not in ['.', '.?', '.<']:
            previousElement = list[-1]
            del list[-1]
            if previousElement in ['.', '.?', '.<']:
                cutoff = previousElement[1] if previousElement in ['.?', '.<'] else ''
                list.append(cutoff + element)
            else:
                list.append(previousElement + ' ' + element)
        else:
            list.append(element)

        return list

    tripleStatements = reduce(repair, trimmed, [''])
    triplesWithNones = list(map(splitIntoTripleParts, tripleStatements))
    triples = [triple for triple in triplesWithNones if triple != None]
    return triples

def splitIntoTripleParts (triple):
    statementPattern = r'(\S+)\s+(\S+)\s+(\S+)'
    statementPatternMatch = re.search(statementPattern, triple)

    if statementPatternMatch:
        return {
            'subject': statementPatternMatch.group(1),
            'predicate': statementPatternMatch.group(2),
            'object': statementPatternMatch.group(3)
        }
    else:
        return None

def fix_URI(query):
	query = re.sub(r"dbr:([^\s]+)" , r"<http://dbpedia.org/resource/\1>" , query)
	if query[-2:]=="}>":
		query = query[:-2]+">}"
	return query