#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Sep 29 23:02:23 2020

@author: noahlefrancois
"""

import numpy as np
import camb
from matplotlib import pyplot as plt

#Provided example function to get the power spectrum from CAMB
def get_spectrum(pars,lmax=2000, fixTau = False):
    if fixTau==False:
        #print('pars are ',pars)
        H0=pars[0]
        ombh2=pars[1]
        omch2=pars[2]
        tau=pars[3]
        As=pars[4]
        ns=pars[5]
    else:
        H0=pars[0]
        ombh2=pars[1]
        omch2=pars[2]
        tau=fixTau
        As=pars[3]
        ns=pars[4]
    pars=camb.CAMBparams()
    pars.set_cosmology(H0=H0,ombh2=ombh2,omch2=omch2,mnu=0.06,omk=0,tau=tau)
    pars.InitPower.set_params(As=As,ns=ns,r=0)
    pars.set_for_lmax(lmax,lens_potential_accuracy=0)
    results=camb.get_results(pars)
    powers=results.get_cmb_power_spectra(pars,CMB_unit='muK')
    cmb=powers['total']
    tt=cmb[:,0]    #you could return the full power spectrum here if you wanted to do say EE
    return tt

"""
Make a small change in each parameter and get the function value at that new location in
order to compute the local gradient. Use central differencing scheme dy/dp = (y_{n+1}-y_{n-1})/(2dp)
"""
def partial_derivatives(x, func, p, dx, tau):
    grad = np.zeros([len(x), len(p)])

    for i in range(len(p)):
        p_dx = p.copy()
        dp = p_dx[i]*dx

        p_dx[i] = p_dx[i] + dp
        y_forward = func(p_dx, fixTau = tau)[2:len(x)+2]
        p_dx[i] = p_dx[i] - 2*dp
        y_backward = func(p_dx, fixTau = tau)[2:len(x)+2]

        dydp_i = (y_forward - y_backward)/(2*dp)
        grad[:,i] = np.asarray(dydp_i)
    return grad

def Newt_Lev_Marq_minimizer(y_data, func, p_guess, err_y, max_iter, tau=False):
    """
    Find the optimal parameters (p) to fit y_data to the model func with fixed tau
    """
    chi_sq = 1e5
    lamb = 1e-4
    dx = 0.01
    grad = partial_derivatives(y_data, func, p_guess, dx, tau)
    chi_store = np.zeros(1)
    p_update = p_guess.copy()
    for i in range(max_iter):
        y_guess = func(p_update, fixTau = tau)[2:len(y_data)+2]
        chi_sq_update = chi_squared(y_data, y_guess, err_y)
        chi_store = np.append(chi_store, [chi_sq_update])
        delta_Chi = chi_sq - chi_sq_update
        if delta_Chi<dx and delta_Chi>0:
            print('Chi-Squared converged to ', chi_sq_update, '\n')
            p_guess = p_update.copy()
            grad = partial_derivatives(y_data, func, p_guess, dx, tau)
            N = np.diag(1/err_y**2)
            cov = np.dot(grad.transpose(),N)
            cov = np.dot(cov,grad)
            cov = np.linalg.inv(cov)
            err = np.sqrt(np.diag(cov))
            break
        elif chi_sq_update > chi_sq:
            print('Chi-Squared increased. Lambda has been increased.')
            lamb = lamb*1e2
            p_update = p_guess.copy()
        else:
            print('Chi-Squared decreased. Lambda has been decreased.')
            chi_sq = chi_sq_update
            p_guess = p_update.copy()
            lamb = lamb*1e-1
            grad = partial_derivatives(y_data, func, p_guess, dx, tau)
        res = np.matrix(y_data - y_guess).transpose()
        grad = np.matrix(grad)
        a = grad.transpose()*np.diag(1/err_y**2)*grad
        A = a + lamb*np.diag(np.diag(a))
        B = grad.transpose()*np.diag(1/err_y**2)*res
        dp= np.linalg.inv(A)*B
        #print(dp)

        for j in range(len(p_guess)):
            p_update[j] = p_update[j] + dp[j]
        print('Iteration ', i, ': \n', 'Chi-Squared = ', chi_sq, '\n Lambda = ', lamb, '\n')
    #N = np.diag(1/err_y**2)
    #cov = np.linalg.pinv(np.dot(np.dot(grad.transpose(),N),grad))
    #err = np.sqrt(np.diag(cov))
    #print('Final fit has Chi-Squared value of ', chi_sq)
    return p_guess, cov, err, chi_store

#Calculate the chi squ value of a fit
def chi_squared(data, fit, error):
    x = np.asarray(data)
    y = np.asarray(fit)
    error = np.asarray(error)
    return sum((x-y)**2/error**2)

plt.ion()

#Load the power spectrum data and store the columns we need
wmap=np.loadtxt('wmap_tt_spectrum_9yr_v5.txt')
multipole = wmap[:,0]; power = wmap[:,1]; errPower = wmap[:,2]

#Calculate the spectrum values using the specified fit params as an initial guess
tau = 0.05
pars_initialGuess_tauFree=np.asarray([65,0.02,0.1,0.05,2e-9,0.96])
pars_initialGuess_tauFixed=np.asarray([65,0.02,0.1,2e-9,0.96])
max_iter = 100

#Fit with tau=0.05 fixed
paramsFix, covFix, errFix, chi_store_Fix = Newt_Lev_Marq_minimizer(power, get_spectrum, pars_initialGuess_tauFixed, errPower, max_iter, tau=0.05)
pFit_Fix = get_spectrum(paramsFix, fixTau = 0.05)[2:len(power)+2]

parameter_labels = ['H_0','O_b h^2','O_c h^2','A_s','Slope']
storeDerivs = np.zeros((len(power),len(paramsFix)))
dx= 0.01

storeDerivs = partial_derivatives(power, get_spectrum, paramsFix, dx, tau)
for i in range(len(paramsFix)):
    plt.semilogy(abs(storeDerivs[:,i]),label='Parameter: {}'.format(parameter_labels[i]))

plt.xlabel('Multipole Moment')
plt.ylabel('Power Spectrum')
plt.title('Numerical Derivatives')
plt.legend(loc='upper right')
"""
This plot of the numerical derivative for each parameter across the data range shows that all of them
stay within a few orders of magnitude, indicating that we should be okay with this dx size and it
should be able to reasonably capture the behaviour of each function.
"""

"""
I expect the errors to be lower when allowing tau to float since more of the parameter space is
accessible to look for a minimum of chi-squared.

Output:

Iteration  0 :
Chi-Squared =  1588.237646582676
Lambda =  1e-05

Chi-Squared decreased. Lambda has been decreased.
Iteration  1 :
 Chi-Squared =  1234.804081089522
 Lambda =  1.0000000000000002e-06

Chi-Squared decreased. Lambda has been decreased.

Iteration  2 :
 Chi-Squared =  1227.936903372273
 Lambda =  1.0000000000000002e-07

Chi-Squared decreased. Lambda has been decreased.

Iteration  3 :
 Chi-Squared =  1227.9356355996376
 Lambda =  1.0000000000000004e-08

Chi-Squared converged to  1227.9356339593712
"""

#Fit with tau as a free parameter
paramsFree, covFree, errFree, chi_store_Free = Newt_Lev_Marq_minimizer(power, get_spectrum, pars_initialGuess_tauFree, errPower, max_iter, tau=False)
pFit_Free = get_spectrum(paramsFree)[2:len(power)+2]

"""
Findings: Chi-Squared is almost the same (1227.9 for fixed, 1227.8 for free) and errors are all
of the same magnitude. However, letting tau be a free parameter results in a much larger number
of iterations before the fit converges. This makes sense since there is an additional degree of
freedom in this fit.

Output:

Chi-Squared decreased. Lambda has been decreased.

Iteration  0 :
 Chi-Squared =  1588.237646582676
 Lambda =  1e-05

Chi-Squared decreased. Lambda has been decreased.

Iteration  1 :
 Chi-Squared =  1235.5192579744582
 Lambda =  1.0000000000000002e-06

Chi-Squared decreased. Lambda has been decreased.

Iteration  2 :
 Chi-Squared =  1227.9836259688445
 Lambda =  1.0000000000000002e-07

Chi-Squared increased. Lambda has been increased.

Iteration  3 :
 Chi-Squared =  1227.9836259688445
 Lambda =  1.0000000000000003e-05

Chi-Squared increased. Lambda has been increased.

Iteration  4 :
 Chi-Squared =  1227.9836259688445
 Lambda =  0.0010000000000000002

Chi-Squared increased. Lambda has been increased.

Iteration  5 :
 Chi-Squared =  1227.9836259688445
 Lambda =  0.10000000000000002

Chi-Squared increased. Lambda has been increased.

Iteration  6 :
 Chi-Squared =  1227.9836259688445
 Lambda =  10.000000000000002

Chi-Squared increased. Lambda has been increased.

Iteration  7 :
 Chi-Squared =  1227.9836259688445
 Lambda =  1000.0000000000002

Chi-Squared decreased. Lambda has been decreased.

Iteration  8 :
 Chi-Squared =  1227.9786348426355
 Lambda =  100.00000000000003

Chi-Squared decreased. Lambda has been decreased.

Iteration  9 :
 Chi-Squared =  1227.972106442598
 Lambda =  10.000000000000004

Chi-Squared decreased. Lambda has been decreased.

Iteration  10 :
 Chi-Squared =  1227.9302550685518
 Lambda =  1.0000000000000004

Chi-Squared decreased. Lambda has been decreased.

Iteration  11 :
 Chi-Squared =  1227.876093263126
 Lambda =  0.10000000000000005

Chi-Squared decreased. Lambda has been decreased.

Iteration  12 :
 Chi-Squared =  1227.8698416859336
 Lambda =  0.010000000000000005

Chi-Squared decreased. Lambda has been decreased.

Iteration  13 :
 Chi-Squared =  1227.8682384894826
 Lambda =  0.0010000000000000007

Chi-Squared decreased. Lambda has been decreased.

Iteration  14 :
 Chi-Squared =  1227.860095673268
 Lambda =  0.00010000000000000007

Iteration  15 :
 Chi-Squared =  1227.8196549427323
 Lambda =  1.0000000000000008e-05

Chi-Squared increased. Lambda has been increased.

Iteration  16 :
 Chi-Squared =  1227.8196549427323
 Lambda =  0.0010000000000000007

Chi-Squared increased. Lambda has been increased.

Iteration  17 :
 Chi-Squared =  1227.8196549427323
 Lambda =  0.10000000000000006

Chi-Squared increased. Lambda has been increased.

Iteration  18 :
 Chi-Squared =  1227.8196549427323
 Lambda =  10.000000000000005

Chi-Squared increased. Lambda has been increased.

Iteration  19 :
 Chi-Squared =  1227.8196549427323
 Lambda =  1000.0000000000006

Chi-Squared increased. Lambda has been increased.

Iteration  20 :
 Chi-Squared =  1227.8196549427323
 Lambda =  100000.00000000006

Chi-Squared increased. Lambda has been increased.

Iteration  21 :
 Chi-Squared =  1227.8196549427323
 Lambda =  10000000.000000006

Chi-Squared increased. Lambda has been increased.

Iteration  22 :
 Chi-Squared =  1227.8196549427323
 Lambda =  1000000000.0000006

Chi-Squared increased. Lambda has been increased.

Iteration  23 :
 Chi-Squared =  1227.8196549427323
 Lambda =  100000000000.00006

Chi-Squared converged to  1227.8196548828296

"""

#Save the results of both fits to a txt file
np.savetxt('Q3_fixedTau_params.txt', paramsFix)
np.savetxt('Q3_fixedTau_err.txt', errFix)
np.savetxt('Q3_fixedTau_cov.txt', covFix)

np.savetxt('Q3_freeTau_params.txt', paramsFree)
np.savetxt('Q3_freeTau_err.txt', errFree)
np.savetxt('Q3_freeTau_cov.txt', covFree)

#Plot the data and the fit, as well as a convergence plot for Chi-Squared

#plt.errorbar(wmap[:,0],wmap[:,1],wmap[:,2],fmt='*')
plt.plot(multipole,power,'.',label='Data')

plt.plot(pFit_Fix,label='Fit, Fixed Tau')
plt.plot(pFit_Free,label='Fit, Free Tau')
plt.xlabel('Multipole Moment')
plt.ylabel('Power Spectrum')
plt.title('WMAP Satellite 9-year CMB Data')
plt.legend()

plt.figure()
plt.plot(chi_store_Fix[1:], label='Fixed Tau')
plt.plot(chi_store_Free[1:], label='Free Tau')
plt.xlabel('Iteration')
plt.ylabel('Chi-Squared')
plt.title('Chi-Squared Convergence')
plt.legend()


from gearbox.migrations import Migration

class AddTableCustClass(Migration):

    database = "common"

    def up(self):
        t = self.table('CustClass', area="Sta_Data_2_256", dump_name="custclas", desc='''
''')
        t.column('CustClass', 'integer', format="9", initial="0", max_width=4, label="Class", column_label="Class", position=2, order=10, help="Customer Class (depend on avg. amount of invoices)")
        t.column('CCName', 'character', format="x(30)", initial="", max_width=60, label="Name", column_label="Name", position=3, order=20, help="Name of the customer class")
        t.column('Amt', 'decimal', format=">>>,>>9.99", decimals=2, initial="0", max_width=17, label="Amt", column_label="Amt", position=4, order=30, help="Lowest limit")
        t.column('Memo', 'character', format="X(60)", initial="", max_width=120, label="Memo", column_label="Memo", position=6, order=50, help="Memo")
        t.column('Brand', 'character', format="x(8)", initial="", max_width=16, label="Brand", column_label="Brand", position=7, order=60, help="Code Of Brand")
        t.index('Class', [['Brand'], ['CustClass'], ['Amt']], area="Sta_Index_2", primary=True, unique=True)
        t.index('Amt', [['Brand'], ['Amt']], area="Sta_Index_2")
        t.index('CCName', [['Brand'], ['CCName'], ['CustClass']], area="Sta_Index_2")

    def down(self):
        self.drop_table('CustClass')


# from aiocoap import Message, NON, Code
# from Bubot.Core.Coap.coap import Message, NON, Code
from urllib.parse import urlparse
from uuid import uuid4

import cbor2

from Bubot.Core.DeviceLink import ResourceLink
from Bubot.Helpers.ExtException import ExtException, dumps_error
from Bubot_CoAP import defines
from Bubot_CoAP.messages.numbers import NON, Code
from Bubot_CoAP.messages.option import Option
from Bubot_CoAP.messages.request import Request
import bson

class OcfMessage:
    def __init__(self, *, fr=None, to=None, ri=None, cn=None, query=None, **kwargs):
        '''
        :param fr: The URI of the message originator.
        :param to: The URI of the recipient of the message.
        :param ri: The  identifier  that  uniquely  identifies  the message in the originator and the recipient.
        :param cn: Information specific to the operation.
        :param query: ?
        :param kwargs: ?
        '''
        self.data = {}
        self.fr = fr
        self.to = to
        self.ri = ri if ri else str(uuid4())
        self.cn = cn  # Content

        # self.uri_path = kwargs.get('uri_path')
        # self.code = kwargs.get('code')
        # self.query = kwargs.get('query', {})
        # self.token = kwargs.get('token')
        # self.mid = kwargs.get('mid')
        # self.data = kwargs.get('data', {})
        # self.raw_msg = None
        # self.code = kwargs.get('code', NON)
        pass

    def parse_url_to(self):
        return urlparse(self.to)

    @staticmethod
    def init_from_bson(data: bytes):
        _data = bson.decode(data)
        if 'op' in _data:
            return OcfRequest(**_data)
        elif 'rs' in _data:
            return OcfResponse(**_data)
        else:
            raise ExtException(message='Bad message data', dump=_data)

    @staticmethod
    def encode_query(query):
        result = []
        for key in query:
            if not isinstance(query[key], list):
                raise Exception('value query param must be only list')
            for elem in query[key]:
                result.append('{}={}'.format(key, elem))
        return result

    @staticmethod
    def decode_query(query):
        result = {}
        if not query:
            return None
        for elem in query:
            key, value = elem.split('=')
            if key not in query:
                result[key] = []
            result[key].append(value)
            pass
        return result

class OcfRequest(OcfMessage):

    def __init__(self, *, obs=None, op=None, multicast=False, **kwargs):
        super().__init__(**kwargs)
        self.obs = obs
        self.op = op
        # self.multicast = multicast
        # if not self.code or not self.op:
        #     if not self.code and self.op:
        #         self.code = self.map_crudn_to_coap_code(self.op)
        #     else:
        #         raise Exception('не заполнены обязательные параметры')

        # self.observe = None

    def to_dict(self):
        return dict(
            fr=self.fr,
            to=self.to,
            ri=self.ri,
            cn=self.cn,
            op=self.op,
            obs=self.obs
        )

    def generate_error(self, err, *, rs=Code.UNPROCESSABLE_ENTITY, obs=None, **kwargs):
        if not isinstance(err, ExtException):
            err = ExtException(parent=err)
        return OcfResponse(
            to=self.fr,
            fr=self.to,
            ri=self.ri,
            cn=err.to_dict(),
            rs=rs,
            obs=obs if obs else self.obs  # todo не понятно как должно быть
            # uri_query=self.query,
            # token=self.token,
            # mid=self.mid,
        )

    def generate_answer(self, data, *, rs=Code.CONTENT, obs=None, **kwargs):
        return OcfResponse(
            to=self.fr,
            fr=self.to,
            ri=self.ri,
            rs=rs,
            cn=data,
            obs=obs if obs else self.obs  # todo не понятно как должно быть
            # mid=kwargs.get('mid', self.mid),
            # uri_query=self.query,
            # token=self.token,
        )
        # if req.raw_msg:
        #     self.raw_msg = req.raw_msg.copy()
        # self.raw_msg.code = Code.CONTENT

    def encode_to_coap(self):

        # try:
        #     parsed = urllib.parse.urlparse(self.to.get_endpoint(), allow_fragments=False)
        #     address = (parsed.hostname, parsed.port)
        # except KeyError:
        #     address = None
        # return msg2, address

        request = Request()
        request.type = NON
        # request.token = generate_random_token(2)
        request.destination = self.to.get_endpoint_address()
        # request.destination = (sender, 5683)
        request.multicast = self.multicast
        request.source = (self.fr.get_endpoint_address()[0], None)
        request.code = self.code
        request.uri_path = self.to.href

        option = Option()
        option.number = defines.OptionRegistry.CONTENT_TYPE.number
        option.value = 10000
        request.add_option(option)

        option = Option()
        option.number = defines.OptionRegistry.URI_QUERY.number
        option.value = self.encode_query(self.query)  # todo  разобраться что тут надо
        request.add_option(option)
        return request, request.destination

        # params = dict(
        #     mtype=NON,
        #     mid=self.mid,
        #     code=self.code,
        #     token=self.token,
        #     uri_host=self.to.anchor,
        #     uri_query=self.encode_query(self.query),
        #     uri_path=self.to.href[1:].split('/'),
        #     # application/vnd.ocf+cbor
        #     # http://www.iana.org/assignments/core-parameters/core-parameters.xhtml#content-formats
        #     content_format=10000,
        #     accept=10000,
        #     observe=self.obs,
        #     payload=cbor2.dumps(self.cn) if self.cn else b''
        # )
        # msg2 = Message(**params)

    @classmethod
    def decode_from_coap(cls, msg, multicast=False):
        # message = Message.decode(raw_data, remote)
        self = cls(
            fr=ResourceLink.init_from_uri('coap://[{0}]:{1}'.format(msg.remote[0], msg.remote[1])),
            # message.opt.uri_path, message.opt.uri_port))]
            to=ResourceLink.init_from_msg(msg),
            code=int(msg.code),
            op=cls.map_coap_code_to_crudn(msg.code.name),
            obs=msg.opt.observe,
            query=cls.decode_query(msg.opt.uri_query),
            token=msg.token,
            mid=msg.mid,
            multicast=multicast
        )
        self.raw_msg = msg
        self.operation = msg.code.name.lower()

        if msg.payload:
            self.cn = cbor2.loads(msg.payload)

        return self

    @staticmethod
    def map_coap_code_to_crudn(code):
        map_coap_to_crudn = {
            'post': 'update',
            'put': 'create',
            'delete': 'delete',
            'get': 'retrieve'
        }
        try:
            return map_coap_to_crudn[code.lower()]
        except KeyError:
            raise Exception('Unknown CRUDN operation ({0})'.format(code))

    @staticmethod
    def map_crudn_to_coap_code(operation):
        #    +------+--------+-----------+
        #    | Code | Name   | Reference |
        #    | 0.01 | GET    | [RFC7252] |
        #    | 0.02 | POST   | [RFC7252] |
        #    | 0.03 | PUT    | [RFC7252] |
        #    | 0.04 | DELETE | [RFC7252] |
        #    +------+--------+-----------+
        map_crudn_to_coap = {
            'create': 3,
            'retrieve': 1,
            'update': 2,
            'delete': 4,
        }
        return map_crudn_to_coap[operation.lower()]

class OcfResponse(OcfMessage):
    def __init__(self, *, rs=None, obs=None, **kwargs):
        super().__init__(**kwargs)
        # self.error = self.is_successful(kwargs['rs'].)
        self.rs = rs  # Response Code RFC 7252
        self.obs = obs  # Observe
        # self.observe = kwargs.get('observe', None)

        # +------+------------------------------+-----------+
        # | Code | Description                  | Reference |
        # +------+------------------------------+-----------+
        # | 2.01 | Created                      | [RFC7252] |
        # | 2.02 | Deleted                      | [RFC7252] |
        # | 2.03 | Valid                        | [RFC7252] |
        # | 2.04 | Changed                      | [RFC7252] |
        # | 2.05 | Content                      | [RFC7252] |
        # | 4.00 | Bad Request                  | [RFC7252] |
        # | 4.01 | Unauthorized                 | [RFC7252] |
        # | 4.02 | Bad Option                   | [RFC7252] |
        # | 4.03 | Forbidden                    | [RFC7252] |
        # | 4.04 | Not Found                    | [RFC7252] |
        # | 4.05 | Method Not Allowed           | [RFC7252] |
        # | 4.06 | Not Acceptable               | [RFC7252] |
        # | 4.12 | Precondition Failed          | [RFC7252] |
        # | 4.13 | Request Entity Too Large     | [RFC7252] |
        # | 4.15 | Unsupported Content - Format | [RFC7252] |
        # | 5.00 | Internal Server Error        | [RFC7252] |
        # | 5.01 | Not Implemented              | [RFC7252] |
        # | 5.02 | Bad Gateway                  | [RFC7252] |
        # | 5.03 | Service Unavailable          | [RFC7252] |
        # | 5.04 | Gateway ExtTimeoutError      | [RFC7252] |
        # | 5.05 | Proxying Not Supported       | [RFC7252] |
        # +------+------------------------------+-----------+

    def to_dict(self):
        return dict(
            fr=self.fr,
            to=self.to,
            ri=self.ri,
            cn=self.cn,
            rs=self.rs,
            obs=self.obs
        )

    def is_successful(self):
        """True if the code is in the successful subrange of the response code range"""
        return True if (64 <= self.rs < 96) else False

    def encode_to_coap(self):
        parsed = urlparse(self.fr.get_endpoint(), allow_fragments=False)
        msg2 = Message(
            mtype=NON,
            mid=self.mid,
            code=self.rs,
            observe=self.obs,
            payload=cbor2.dumps(self.cn) if self.cn else b'',
            token=self.token,
            uri_query=self.encode_query(self.query),
            uri_host=self.to.anchor,
            uri_path=self.to.href[1:].split('/'),
            content_format=10000,
            accept=10000
            # iotivity_addresses=b'\xc0'
        )
        # message = self.raw_msg
        # message.code = self.code
        # message.payload = cbor2.dumps(self.data) if self.data else b''
        # message.mtype = NON
        return msg2, (parsed.hostname, parsed.port)

    @classmethod
    def decode_from_coap(cls, msg, multicast=False):
        # message = Message.decode(raw_data, remote)
        self = cls(
            code=int(msg.code),
            rs=msg.code.name.lower(),
            token=msg.token,
            mid=msg.mid,
            query=cls.decode_query(msg.opt.uri_query),
            obs=msg.opt.observe,
            multicast=multicast,
            to=ResourceLink.init_from_msg(msg)
        )
        self.raw_msg = msg
        self.operation = msg.code.name.lower()
        if msg.payload:
            self.cn = cbor2.loads(msg.payload)
        if msg.opt.uri_query:
            for elem in msg.opt.uri_query:
                key, value = elem.split('=')
                if key not in self.uri_query:
                    self.uri_query[key] = []
                self.uri_query[key].append(value)
                pass
        return self

    # def to_dict(self):
    #     return dict(
    #         code=self.code,
    #         token=self.token,
    #         mid=self.mid,
    #         obs=self.obs,
    #         fr=self.fr.data if self.fr else None,
    #         to=self.to.data if self.to else None,
    #         cn=self.cn
    #     )
    #
    # def __str__(self):
    #     return '{code} {value}'.format(value=str(self.cn))


import numpy as np
from torch.utils.data import DataLoader, Dataset, random_split, ConcatDataset
import torch
from typing import Tuple, List

class DIADataSet(Dataset):
    def __init__(
        self,
        libMsMz: np.ndarray,
        libMs: np.ndarray,
        libSpectrumIfShare: np.ndarray,
        libTarget: np.ndarray,
    ) -> None:
        r"""
        Input:
        * libMsMz 肽段匹配到的 s 个最相关的混合图谱
        * libMs 肽段理论二级图谱
        * libSpectrumIfShare 肽段特异峰向量
        * libTarget 标签
        """
        super(DIADataSet, self).__init__()
        self.libSpectrumIfShare = libSpectrumIfShare
        self.libMsMz = libMsMz
        self.libMs = libMs
        self.libTarget = libTarget
        self.len = len(self.libTarget)

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

    def __getitem__(
        self, index: int
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        r""" """
        return (
            self.libMsMz[index],
            self.libMs[index],
            self.libSpectrumIfShare[index],
            self.libTarget[index],
        )

def SplitTrainAndValDataSet(
    TrainDataSet: DIADataSet,
    TrainDataDecoySet: DIADataSet,
    lengthValDataSet: int,
    lengthValDataDecoySet: int,
) -> Tuple[DIADataSet, DIADataSet, DIADataSet]:
    r"""
    Input:
    * TrainDataSet: 正样本训练数据
    * TrainDataDecoySet 负样本训练数据
    * lengthValDataSet 正样本验证集大小
    * lengthValDataDecoySet 负样本验证集大小
    1. 根据输入的训练数据, 从正样本和负样本中各抽取数条数据作为验证集，每个 epoch 结束前进行 FDR 控制
       选取大于分数阈值最多的模型
    2. 将剩余训练数据中的正样本和负样本进行合并并且进行随机化选取

    Output:
        返回训练集, 验证集, 诱饵验证集
    """
    trainDataSet, valDataSet = random_split(
        TrainDataSet,
        [len(TrainDataSet) - lengthValDataSet, lengthValDataSet],
        generator=torch.Generator().manual_seed(0),
    )
    trainDataDecoySet, valDataDecoySet = random_split(
        TrainDataDecoySet,
        [len(TrainDataDecoySet) - lengthValDataDecoySet, lengthValDataDecoySet],
        generator=torch.Generator().manual_seed(0),
    )
    trainDataSet = ConcatDataset([trainDataSet, trainDataDecoySet])
    return trainDataSet, valDataSet, valDataDecoySet

def GetAllDataLoader(
    batch_size: int,
    TrainDataSet: DIADataSet,
    TrainDataSetDecoy: DIADataSet,
    testDataSet: DIADataSet,
    testDataSetDecoy: DIADataSet,
    lengthValDataSet: int,
    lengthValDataSetDecoy: int,
) -> Tuple[DataLoader, DataLoader, DataLoader, DataLoader, DataLoader]:
    r"""
    把 GetTrainDataLoader 和 GetTestDataLoader 两个函数操作合并在一起
    """
    trainDataSet, valDataSet, valDataDecoySet = SplitTrainAndValDataSet(
        TrainDataSet, TrainDataSetDecoy, lengthValDataSet, lengthValDataSetDecoy
    )
    trainData, valData, valDataDecoy = GetTrainDataLoader(
        batch_size, trainDataSet, valDataSet, valDataDecoySet
    )
    testData, testDataDecoy = GetTestDataLoader(
        batch_size, testDataSet, testDataSetDecoy
    )
    return trainData, valData, valDataDecoy, testData, testDataDecoy

def GetTrainDataLoader(
    batch_size: int,
    trainDataSet: DIADataSet,
    valDataSet: List[DIADataSet],
    valDataDecoySet: List[DIADataSet],
) -> Tuple[DataLoader, List[DataLoader], List[DataLoader]]:
    r"""
    Input:
    * batch_size: batch 大小
    * trainDataSet: 训练集
    * valDataSet: 验证集
    * valDataDecoySet: 诱饵验证集

    OutPut: return
    * 训练集
    * 正样本验证集
    * 诱饵验证集的 dataloader
    """
    trainData = DataLoader(
        dataset=trainDataSet, shuffle=True, batch_size=batch_size, num_workers=8
    )
    valData = [
        DataLoader(dataset=dataset, shuffle=False,
                   batch_size=batch_size, num_workers=8)
        for dataset in valDataSet
    ]
    valDataDecoy = [
        DataLoader(dataset=dataset, shuffle=False,
                   batch_size=batch_size, num_workers=8)
        for dataset in valDataDecoySet
    ]
    return trainData, valData, valDataDecoy

def GetTestDataLoader(
    batch_size: int, testDataSet: DIADataSet, testDataDecoySet: DIADataSet
) -> Tuple[DataLoader, DataLoader]:
    r"""
    Input:
    * batch_size: batch 大小
    * testDataSet: 测试集
    * testDataDecoySet: 诱饵测试集

    Output: return
    * 测试集
    * 诱饵测试集的 dataloader
    """
    testData = DataLoader(
        dataset=testDataSet, shuffle=False, batch_size=batch_size, num_workers=8
    )
    testDataDecoy = DataLoader(
        dataset=testDataDecoySet, shuffle=False, batch_size=batch_size, num_workers=8
    )
    return testData, testDataDecoy

def LibSpectrumIfSharePreTreatMent(LibSpectrumIfShare: np.ndarray) -> np.ndarray:
    r"""
    对特异峰向量进行预处理操作
    [0, 0, 0, 0, 1, 1] -> [[0], [0], [0], [0], [1], [1]]
    """
    x = LibSpectrumIfShare
    y = []
    for item in x:
        y.append(item)
    testLibSpectrumIfShare = np.array(y)
    testLibSpectrumIfShare = np.reshape(
        testLibSpectrumIfShare, (np.shape(testLibSpectrumIfShare)[0], 6, 1)
    )
    return testLibSpectrumIfShare

def LibMsMzPreTreatMent(LibMsMz: np.ndarray) -> np.ndarray:
    r"""
    1. 匹配到的谱图进行预处理操作
    2. 归一化等或不进行操作
    """
    Ms = []
    for data in LibMsMz:
        Ms.append(data)
    test = np.array(Ms)
    msIntensity = test[:, :, :, [1]]
    msIntensity = msIntensity
    msMz = test[:, :, :, [0]]
    Ms = np.concatenate((msMz, msIntensity), axis=3)
    LibMsMz = Ms
    return LibMsMz

def Pretreatment(
    data: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    r"""
    对图谱和特异峰向量进行预处理操作
    """
    LibMsMz, LibMz, LibTarget, LibSpectrumIfShare, fileWinLib = (
        data[0],
        data[1],
        data[2],
        data[3],
        data[4],
    )
    LibSpectrumIfShare = LibSpectrumIfSharePreTreatMent(LibSpectrumIfShare)
    LibMsMz = LibMsMzPreTreatMent(LibMsMz)
    return LibMsMz, LibMz, LibTarget, LibSpectrumIfShare, fileWinLib

def GetTrainDataSet(data: np.ndarray) -> DIADataSet:
    r"""
    Input:
    * data 输入的初始训练数据

    事先对肽段匹配上的图谱、肽段参考图谱、肽段特异峰向量进行预处理操作, 详情见 Pretreatment 函数

    Output: return
    * 处理之后的数据集
    """
    LibMsMz, LibMz, LibTarget, LibSpectrumIfShare, _ = Pretreatment(data)
    return DIADataSet(LibMsMz, LibMz, LibSpectrumIfShare, LibTarget)

def GetTestDataSetAndFilewindow(
    data: np.ndarray,
) -> Tuple[DIADataSet, Tuple[str, str, Tuple[str, int]]]:
    r"""
    Input:
    * data 输入的初始测试数据

    事先对肽段匹配上的图谱、肽段参考图谱、肽段特异峰向量进行预处理操作, 详情见 Pretreatment 函数

    Output: return
    * 处理之后的数据集
    * (文件名、窗口信息、肽段名)
    """
    LibMsMz, LibMz, LibTarget, LibSpectrumIfShare, FileWinLib = Pretreatment(
        data)
    return DIADataSet(LibMsMz, LibMz, LibSpectrumIfShare, LibTarget), FileWinLib

def GetTestData(
    batch_size: int, testPath: str, testDecoyPath: str
) -> Tuple[DataLoader, DataLoader, Tuple[str], Tuple[str]]:
    r"""
    Input:
    * batch_size batch 大小
    * testPath 测试数据路径
    * testDecoyPath 测试诱饵数据路径

    Output: return
    * 测试数据 DataLoader
    * 测试诱饵数据 DataLoader
    * 测试数据 (文件, 窗口, (肽段名, 带电荷数))
    * 测试诱饵数据 (文件, 窗口, (肽段名, 带电荷数))
    """
    TestData = np.load(testPath, allow_pickle=True)
    TestDataDecoy = np.load(testDecoyPath, allow_pickle=True)

    testDataSet, fileWinLib = GetTestDataSetAndFilewindow(TestData)
    testDataDecoySet, fileWinLibDecoy = GetTestDataSetAndFilewindow(
        TestDataDecoy)

    testData, testDataDecoy = GetTestDataLoader(
        batch_size, testDataSet, testDataDecoySet
    )

    return testData, testDataDecoy, fileWinLib, fileWinLibDecoy

def GetLFQAndTrypsin_HFXData(
    batch_size: int,
    lfqPath: str,
    lfqDecoyPath: str,
    Trypsin_HFXpath: str,
    Trypsin_HFXdecoyPath: str,
) -> Tuple[DataLoader, List[DataLoader], List[DataLoader]]:
    r"""
    Input:
    * batch_size: batch 大小
    * lfqPath: lfq 正样本训练数据路径  (可能为 lfq32 和 lfq64)
    * lfqDecoyPath: lfq 负样本训练数据路径 (可能为 lfq32 和 lfq64)
    * Trypsin_HFXPath: 血浆正样本训练数据路径
    * Trypsin_HFXDecoyPath: 血浆负样本训练数据路径

    获取 lfq 和 trupsin_HFX 血浆的训练数据

    Output: return
    * 训练集
    * 正样本验证集
    * 诱饵验证集的 dataloader
    """
    # lfq 和血浆数据
    lfqTrainData = np.load(lfqPath, allow_pickle=True)
    lfqDecoyData = np.load(lfqDecoyPath, allow_pickle=True)
    Trypsin_HFXTrainData = np.load(Trypsin_HFXpath, allow_pickle=True)
    Trypsin_HFXDecoyData = np.load(Trypsin_HFXdecoyPath, allow_pickle=True)
    # 获取数据的 DataSet
    lfqTrainDataSet = GetTrainDataSet(lfqTrainData)
    lfqTrainDataDecoySet = GetTrainDataSet(lfqDecoyData)
    Trypsin_HFXTrainDataSet = GetTrainDataSet(Trypsin_HFXTrainData)
    Trypsin_HFXTrainDataDecoySet = GetTrainDataSet(Trypsin_HFXDecoyData)

    """
    由于血浆和 lfq 样本数量不一致, 有训练集和验证集拆分可以采用两种不同的方法
    一、
        实现将两者进行合并处理, 再从正样本和负样本中各抽取 2w 条数据作为验证集和诱饵验证集

    二、
    1. 血浆从正样本和负样本中各抽取 5k 条作为验证集和诱饵验证集, lfq 从正样本和负样本中抽取 1w 条作为验证集和诱饵验证集
    2. 把两者训练集、验证集、诱饵验证集进行合并, 训练集随机化
    """
    # TrainDataSet = ConcatDataset([lfqTrainDataSet, Trypsin_HFXTrainDataSet])
    # TrainDataDecoySet = ConcatDataset(
    #     [lfqTrainDataDecoySet, Trypsin_HFXTrainDataDecoySet]
    # )

    # trainDataSet, valDataSet, valDataDecoySet = SplitTrainAndValDataSet(
    #     TrainDataSet, TrainDataDecoySet, 20000, 20000
    # )

    lfqtrainDataSet, lfqvalDataSet, lfqvalDataDecoySet = SplitTrainAndValDataSet(
        lfqTrainDataSet, lfqTrainDataDecoySet, 10000, 10000
    )

    (
        Trypsin_HFXtrainDataSet,
        Trypsin_HFXvalDataSet,
        Trypsin_HFXvalDataDecoySet,
    ) = SplitTrainAndValDataSet(
        Trypsin_HFXTrainDataSet, Trypsin_HFXTrainDataDecoySet, 1000, 1000
    )

    trainDataSet = ConcatDataset([lfqtrainDataSet, Trypsin_HFXtrainDataSet])
    # ConcatDataset([lfqvalDataSet, Trypsin_HFXvalDataSet]),
    # ConcatDataset([lfqvalDataDecoySet, Trypsin_HFXvalDataDecoySet]),

    valDataSet, valDataDecoySet = [lfqvalDataSet, Trypsin_HFXvalDataSet], [
        lfqvalDataDecoySet,
        Trypsin_HFXvalDataDecoySet,
    ]

    trainData, valData, valDataDecoy = GetTrainDataLoader(
        batch_size, trainDataSet, valDataSet, valDataDecoySet
    )
    return trainData, valData, valDataDecoy

def GetAllTrainData(batch_size: int,
                    lfq32Path: str,
                    lfq32DecoyPath: str,
                    lfq64Path: str,
                    lfq64DecoyPath: str,
                    tims20211002Path: str,
                    tims20211002DecoyPath: str,
                    Trypsin_HFXpath: str,
                    Trypsin_HFXdecoyPath: str) -> Tuple[DataLoader, List[DataLoader], List[DataLoader]]:
    r"""
    Input:
    * batch_size: batch 大小
    * lfq32Path: lfq32 正样本训练数据路径
    * lfq32DecoyPath: lfq32 负样本训练数据路径
    * lfq64Path: lfq64 正样本训练数据路径
    * lfq64DecoyPath: lfq64 负样本训练数据路径
    * tims20211002Path: tims 血浆正样本训练数据路径
    * tims20211002Path: tims 血浆负样本训练数据路径
    * Trypsin_HFXPath: Trypsin_HFX 血浆正样本训练数据路径
    * Trypsin_HFXDecoyPath: Trypsin_HFX 血浆负样本训练数据路径

    获取 lfq 和两类血浆的训练数据

    Output: return
    * 训练集
    * 正样本验证集
    * 诱饵验证集的 dataloader
    """
    # lfq 和血浆数据
    lfq32TrainData = np.load(lfq32Path, allow_pickle=True)
    lfq32DecoyData = np.load(lfq32DecoyPath, allow_pickle=True)
    lfq64TrainData = np.load(lfq64Path, allow_pickle=True)
    lfq64DecoyData = np.load(lfq64DecoyPath, allow_pickle=True)
    tims20211002TrainData = np.load(tims20211002Path, allow_pickle=True)
    tims20211002DecoyData = np.load(tims20211002DecoyPath, allow_pickle=True)
    Trypsin_HFXTrainData = np.load(Trypsin_HFXpath, allow_pickle=True)
    Trypsin_HFXDecoyData = np.load(Trypsin_HFXdecoyPath, allow_pickle=True)
    # 获取数据的 DataSet
    lfq32TrainDataSet = GetTrainDataSet(lfq32TrainData)
    lfq32TrainDataDecoySet = GetTrainDataSet(lfq32DecoyData)
    lfq64TrainDataSet = GetTrainDataSet(lfq64TrainData)
    lfq64TrainDataDecoySet = GetTrainDataSet(lfq64DecoyData)
    tims20211002TrainDataSet = GetTrainDataSet(tims20211002TrainData)
    tims20211002TrainDataDecoySet = GetTrainDataSet(tims20211002DecoyData)
    Trypsin_HFXTrainDataSet = GetTrainDataSet(Trypsin_HFXTrainData)
    Trypsin_HFXTrainDataDecoySet = GetTrainDataSet(Trypsin_HFXDecoyData)

    lfq32trainDataSet, lfq32valDataSet, lfq32valDataDecoySet = SplitTrainAndValDataSet(
        lfq32TrainDataSet, lfq32TrainDataDecoySet, 10000, 10000
    )

    lfq64trainDataSet, lfq64valDataSet, lfq64valDataDecoySet = SplitTrainAndValDataSet(
        lfq64TrainDataSet, lfq64TrainDataDecoySet, 10000, 10000
    )

    tims20211002trainDataSet, tims20211002valDataSet, tims20211002valDataDecoySet = SplitTrainAndValDataSet(
        tims20211002TrainDataSet, lfq64TrainDataDecoySet, 1000, 1000
    )

    (
        Trypsin_HFXtrainDataSet,
        Trypsin_HFXvalDataSet,
        Trypsin_HFXvalDataDecoySet,
    ) = SplitTrainAndValDataSet(
        Trypsin_HFXTrainDataSet, Trypsin_HFXTrainDataDecoySet, 1000, 1000
    )

    trainDataSet = ConcatDataset(
        [lfq32trainDataSet, lfq64trainDataSet, tims20211002trainDataSet, Trypsin_HFXtrainDataSet])

    valDataSet, valDataDecoySet = [lfq32valDataSet, lfq64valDataDecoySet, tims20211002valDataSet, Trypsin_HFXvalDataSet], [
        lfq32valDataDecoySet,
        lfq64valDataDecoySet,
        tims20211002valDataSet,
        Trypsin_HFXvalDataDecoySet,
    ]

    trainData, valData, valDataDecoy = GetTrainDataLoader(
        batch_size, trainDataSet, valDataSet, valDataDecoySet
    )
    return trainData, valData, valDataDecoy

def GetTrainDataPath() -> Tuple[str, str, str, str, str, str, str, str]:
    r"""
    获取各训练数据路径, 包含以下各类数据
    * LFQ32fix
    * LFQ64var
    * Trypsin_HFX 血浆
    * tims 血浆
    """
    root = "/data/xp/SpectronautLabelTrainAndTest/"

    lfq32TrainPath = (
        root
        + "LFQ32/trainData/HYE124_TTOF6600_32fix_train_win_data_minmz5_identify.npy"
    )

    lfq32TrainDecoyPath = (
        root
        + "LFQ32/trainData/HYE124_TTOF6600_32fix_train_win_data_minmz5_identify_decoy.npy"
    )

    lfq64TrainPath = (
        root
        + "LFQ64/trainData/HYE124_TTOF6600_64var_train_win_data_minmz5_identify.npy"
    )

    lfq64TrainDecoyPath = (
        root
        + "LFQ64/trainData/HYE124_TTOF6600_64var_train_win_data_minmz5_identify_decoy.npy"
    )

    Trypsin_HFXtrainpath = (
        root
        + "20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX/trainData/20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX_train_win_data_minmz5_identify.npy"
    )

    Trypsin_HFXtrainDecoyPath = (
        root
        + "20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX/trainData/20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX_train_win_data_minmz5_identify_decoy.npy"
    )

    tims20211002TrainPath = (
        root
        + "tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc/trainData/tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc_train_win_data_minmz5_identify.npy"
    )

    tims20211002TrainDecoyPath = (
        root
        + "tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc/trainData/tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc_train_win_data_minmz5_identify_decoy.npy"
    )

    return (
        lfq32TrainPath,
        lfq32TrainDecoyPath,
        lfq64TrainPath,
        lfq64TrainDecoyPath,
        Trypsin_HFXtrainpath,
        Trypsin_HFXtrainDecoyPath,
        tims20211002TrainPath,
        tims20211002TrainDecoyPath,
    )

def GetTestDataPath() -> (
    Tuple[str, str, str, str, str, str, str, str, str, str, str, str]
):
    r"""
    获取各训练数据路径, 包含以下各类数据
    * LFQ32fix
    * LFQ64var
    * Trypsin_HFX 血浆
    * tims 血浆
    * Hela 数据
    * 无色谱 lfq
    """
    root = "/data/xp/SpectronautLabelTrainAndTest/"

    lfq32TestPath = (
        root + "LFQ32/testData/HYE124_TTOF6600_32fix_test_win_data_minmz5_identify.npy"
    )

    lfq32TestDecoyPath = (
        root
        + "LFQ32/testData/HYE124_TTOF6600_32fix_test_win_data_minmz5_identify_decoy.npy"
    )

    lfq64TestPath = (
        root + "LFQ64/testData/HYE124_TTOF6600_64var_test_win_data_minmz5_identify.npy"
    )

    lfq64TestDecoyPath = (
        root
        + "LFQ64/testData/HYE124_TTOF6600_64var_test_win_data_minmz5_identify_decoy.npy"
    )

    Trypsin_HFXtestpath = (
        root
        + "20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX/testData/20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX_test_win_data_minmz5_identify.npy"
    )

    Trypsin_HFXtestDecoyPath = (
        root
        + "20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX/testData/20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX_test_win_data_minmz5_identify_decoy.npy"
    )

    tims20211002TestPath = (
        root
        + "tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc/testData/tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc_test_win_data_minmz5_identify.npy"
    )

    tims20211002TestDecoyPath = (
        root
        + "tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc/testData/tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc_test_win_data_minmz5_identify_decoy.npy"
    )

    HelaTestPath = root + "Hela/testData/Hela_test_win_data_minmz5_identify.npy"

    HelaTestDecoyPath = (
        root + "Hela/testData/Hela_test_win_data_minmz5_identify_decoy.npy"
    )

    NoRTtestPath = root + "NoRTLfq/testData/NoRT_LFQ_test_win_data_minmz5_identify.npy"

    NoRTtestDecoyPath = (
        root + "NoRTLfq/testData/NoRT_LFQ_test_win_data_minmz5_identify_decoy.npy"
    )

    return (
        lfq32TestPath,
        lfq32TestDecoyPath,
        lfq64TestPath,
        lfq64TestDecoyPath,
        Trypsin_HFXtestpath,
        Trypsin_HFXtestDecoyPath,
        tims20211002TestPath,
        tims20211002TestDecoyPath,
        HelaTestPath,
        HelaTestDecoyPath,
        NoRTtestPath,
        NoRTtestDecoyPath,
    )

def GetSpectronautResultPath() -> Tuple[str, str, str, str]:
    r"""
    获取 Spectronaut 结果的路径
    """
    root = "/data/xp/SpectronautLabelTrainAndTest/"
    lfq32Path = root + "LFQ32/LFQ32_Spectronaut_HYE124_TTOF6600_32fix_identify_all_pep.npy"
    lfq64Path = root + "LFQ64/LFQ64_Spectronaut_HYE124_TTOF6600_64var_quantify_all_pep.npy"
    tims20211002Path = root + \
        "tims-TOF_20211002_30min_LCH_MN_144-Plasma_SPEED-DIA_300ng_mix-qc/plasma_1_2_3_4_6_7_8_9_Spectronaut_identify_all_pep.npy"
    Trypsin_HFXPath = root + \
        "20230111_MN_Plasma_DIA_formal_sample_trypsin_HFX/20230111_plasma_24QC_Spectronaut_identify_all_pep.npy"
    return lfq32Path, lfq64Path, tims20211002Path, Trypsin_HFXPath


from lesson1 import locate_card, count_rot_linear, count_rotations
import pytest

# query is the first element
testdata1 ={
    'input': {
        'cards': [4, 2, 1, -1],
        'query': 4
    },
}

# print(locate_card(testdata1['input']['cards'], testdata1['input']['query']))

def test_1() -> None:
    assert locate_card(testdata1['input']['cards'], testdata1['input']['query']) == 0

# below are tests suggested by course

# store tests as list
tests = []

# query occurs in the middle
tests.append({
    'input': {
        'cards': [13, 11, 10, 7, 4, 3, 1, 0],
        'query': 1
    },
    'output': 6
})
# query is the first element
tests.append({
    'input': {
        'cards': [4, 2, 1, -1],
        'query': 4
    },
    'output': 0
})
# query is the last element
tests.append({
    'input': {
        'cards': [3, -1, -9, -127],
        'query': -127
    },
    'output': 3
})
# cards contains just one element, query
tests.append({
    'input': {
        'cards': [6],
        'query': 6
    },
    'output': 0
})

# cards does not contain query, assuming then that -1 is expected output
tests.append({
    'input': {
        'cards': [9, 7, 5, 2, -9],
        'query': 4
    },
    'output': -1
})

# cards is empty
tests.append({
    'input': {
        'cards': [],
        'query': 7
    },
    'output': -1
})
# numbers can repeat in cards
tests.append({
    'input': {
        'cards': [8, 8, 6, 6, 6, 6, 6, 3, 2, 2, 2, 0, 0, 0],
        'query': 3
    },
    'output': 7
})

# query occurs multiple times
tests.append({
    'input': {
        'cards': [8, 8, 6, 6, 6, 6, 6, 6, 3, 2, 2, 2, 0, 0, 0],
        'query': 6
    },
    'output': 2
})

# rotated list tests

testsrotate=[]

# rotated once
testsrotate.append({
    'input': {
        'cards': [5, 1, 2, 3, 4],
    },
    'output': 1
})

# A list that wasn't rotated at all.
testsrotate.append({
    'input': {
        'cards': [1, 2, 3, 4, 5, 6, 7, 8],
    },
    'output': 0
})

# A list that was rotated n-1 times, where n is the size of the list
testsrotate.append({
    'input': {
        'cards': [2, 3, 4, 5, 6, 7, 8, 1],
    },
    'output': 7
})

# A list containing just one element.
testsrotate.append({
    'input': {
        'cards': [1],
    },
    'output': 0
})

# tests for problem 1
@pytest.mark.parametrize("input_data", tests)
def test_locate_card(input_data)-> None:
    cards = input_data['input']['cards']
    query = input_data['input']['query']
    expected_output = input_data['output']
    assert locate_card(cards, query) == expected_output

# tests for rotated list
@pytest.mark.parametrize("input_data", testsrotate)
def test_rotate_card(input_data)-> None:
    nums = input_data['input']['cards']
    expected_output = input_data['output']
    assert count_rot_linear(nums) == expected_output

@pytest.mark.parametrize("input_data", testsrotate)
def test_binary_rotate(input_data)-> None:
    nums = input_data['input']['cards']
    expected_output = input_data['output']
    assert count_rotations(nums) == expected_output

Lesson 1 - Binary search, linked lists, complexity/test_lesson1.py


# -*- coding: utf-8 -*-

"""
Created on Fri May 6 7:11:32 2016
@author: mdclarke
mnefun processing script for face to face
"""
import mnefun
import numpy as np
from score import score

n_cores = 4
params = mnefun.Params(tmin=-0.1, tmax=7.0, n_jobs=n_cores,
                       n_jobs_mkl=1, lp_cut=80., lp_trans=3.,
                       n_jobs_fir=n_cores, n_jobs_resample=n_cores,
                       bmin=-0.1, decim=10, proj_sfreq=200,
                       filter_length='auto')
# Notes
# 108 has 19s for 13s

params.subjects = ['f2f_009_01_OTP', 'f2f_010_01_OTP', 'f2f_011_01_OTP',
                   'f2f_012_01_OTP', 'f2f_013_01_OTP', 'f2f_014_01_OTP',
                   'f2f_015_01_OTP', 'f2f_018_01_OTP', 'f2f_021_01_OTP',
                   'f2f_022_01_OTP', 'f2f_023_01_OTP', 'f2f_024_01_OTP',
                   'f2f_025_01_OTP', 'f2f_026_01_OTP', 'f2f_027_01_OTP',
                   'f2f_028_01_OTP', 'f2f_029_01_OTP', 'f2f_034_01_OTP',
                   'f2f_036_01_OTP', 'f2f_037_01_OTP', 'f2f_038_01_OTP',
                   'f2f_102_OTP', 'f2f_103_OTP', 'f2f_104_OTP',
                   'f2f_106_OTP', 'f2f_107_OTP', 'f2f_108_OTP',
                   'f2f_111_OTP', 'f2f_112_OTP', 'f2f_114_OTP',
                   'f2f_115_OTP', 'f2f_116_OTP', 'f2f_117_01_OTP',
                   'f2f_118_01_OTP', 'f2f_119_01_OTP',
                   'f2f_013_sim']

params.subjects_dir = 'subjects'
params.structurals = params.subjects
params.dates = [(2013, 0, 00)] * len(params.subjects)
params.subject_indices = [35]
# params.subject_indices = np.setdiff1d(np.arange(len(params.subjects)),
#                                       [])
params.score = score
params.acq_ssh = 'maggie@minea.ilabs.uw.edu'
params.acq_dir = ['/sinuhe/data03/face_to_face']
params.sws_ssh = 'localhost'  # 'maggie@kasga.ilabs.uw.edu'
params.sws_dir = '/mnt/bakraid/data/sss_work'  # '/data07/maggie/f2f/'
params.sws_port = 2222
# SSS options
params.sss_type = 'python'
params.sss_regularize = 'in'
params.tsss_dur = 4.
params.int_order = 6
params.st_correlation = .995  # otp
params.movecomp = 'inter'
params.trans_to = 'twa'
params.coil_bad_count_duration_limit = np.inf
params.coil_dist_limit = None
params.coil_t_window = 'auto'
params.coil_t_step_min = 1.
# Trial/CH rejection criteria
params.autoreject_thresholds = True  # set automatically for each subject
params.autoreject_types = ('mag', 'grad')
params.auto_bad_reject = None
params.ssp_ecg_reject = None
params.flat = dict(grad=1e-13, mag=1e-15)

params.compute_rank = True  # better inverse
params.cov_rank = None  # compute cov on SSS basis
params.force_erm_cov_rank_full = False
params.proj_ave = True  # compute projs on Evoked artifact waveform
params.run_names = ['%s']
params.get_projs_from = np.arange(1)
params.inv_names = ['%s']
params.inv_runs = [np.arange(1)]
params.runs_empty = ['%s_erm']
params.proj_nums = [[2, 2, 0],  # ECG: grad/mag/eeg
                    [0, 0, 0],  # EOG
                    [0, 0, 0]]  # Continuous (from ERM)
params.cov_method = 'shrunk'
params.bem_type = '5120'
# Epoching
params.in_names = ['talk', 'ignore']

params.in_numbers = [13, 15]

params.analyses = ['All',
                   'Each']

params.out_names = [['All'],
                    ['talk', 'ignore']]

params.out_numbers = [[1, 1],  # Combine all trials
                      [1, 2]]  # Each
params.must_match = [
    [],  # no match
    [0, 1]]  # match each

times = [0.5, 1., 2.5]

params.report_params.update(  # add a couple of nice diagnostic plots
    bem=False,  # Using a surrogate
    good_hpi_count=False,
    whitening=[
        dict(analysis='All', name='aud',
             cov='%s-80-sss-cov.fif'),
        dict(analysis='Each', name='talk',
             cov='%s-80-sss-cov.fif'),
        dict(analysis='Each', name='ignore',
             cov='%s-80-sss-cov.fif'),
    ],
    sensor=False,
    source=[
        dict(analysis='All', name='All',
             inv='%s_aud-80-sss-meg-free-inv.fif', times=times,
             views=['lat', 'caudal'], size=(800, 800)),
        dict(analysis='Each', name='talk',
             inv='%s_aud-80-sss-meg-free-inv.fif', times=times,
             views=['lat', 'caudal'], size=(800, 800)),
        dict(analysis='Each', name='ignore',
             inv='%s_vis-80-sss-meg-free-inv.fif', times=times,
             views=['lat', 'caudal'], size=(800, 800)),
    ],
    source_alignment=False,
    psd=False,  # often slow
)
mnefun.do_processing(
    params,
    fetch_raw=False,
    do_score=True,
    push_raw=False,
    do_sss=True,
    fetch_sss=True,
    do_ch_fix=False,
    gen_ssp=True,
    apply_ssp=True,
    write_epochs=True,
    gen_covs=True,
    gen_fwd=True,
    gen_inv=True,
    gen_report=True,
    print_status=False)


# -*- coding: utf-8 -*-

import arcpy
import math
from time import sleep

class TPGeometry(arcpy.Geometry):
    """Equivalent to ArcPy Geometry with some extra methods and properties."""
    secondPoint = None
    beforeLastPoint = None

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

        self._getSecondPoints()

    def _getSecondPoints(self):
        """Finds the secondPoint and the beforeLastPoint."""
        pt_array = self.getPart()[0]

        self.secondPoint = pt_array[1]
        self.beforeLastPoint = pt_array[-2]

    def invert(self):
        """
        Inverts the position of the vertices in the geometry.

        :return: inverted geometry
        """
        pt_array = self.getPart()[0]
        inverted_array = arcpy.Array(list(reversed(pt_array)))
        inverted_geom = TPGeometry(self.type, inverted_array)

        return inverted_geom

    def bearing(self, start=None, end=None):
        """
        Finds the azimuth of a line.

        :param start: the start of the vector (default = firstPoint)
        :param end: the end of the vector (default = lastPoint)

        :return: the azimuth in degrees
        """

        if start is None: start = self.firstPoint
        if end is None: end = self.lastPoint

        bearing = None

        xa = start.X
        ya = start.Y
        xb = end.X
        yb = end.Y

        if xa < xb and ya < yb:
            # quadrant I
            a = xb - xa
            b = yb - ya
            c = math.sqrt(a ** 2 + b ** 2)
            bearing = math.degrees(math.asin(a / c))

        elif xa < xb and ya > yb:
            # quadrant II
            a = ya - yb
            b = xb - xa
            c = math.sqrt(a ** 2 + b ** 2)
            bearing = 90 + math.degrees(math.asin(a / c))

        elif xa > xb and ya > yb:
            # quadrant III
            a = xa - xb
            b = ya - yb
            c = math.sqrt(a ** 2 + b ** 2)
            bearing = 180 + math.degrees(math.asin(a / c))

        elif xa > xb and ya < yb:
            # quadrant IV
            a = yb - ya
            b = xa - xb
            c = math.sqrt(a ** 2 + b ** 2)
            bearing = 270 + math.degrees(math.asin(a / c))

        elif xa == xb and ya < yb:
            bearing = 0

        elif xa < xb and ya == yb:
            bearing = 90

        elif xa == xb and ya > yb:
            bearing = 180

        elif xa > xb and ya == yb:
            bearing = 270

        if bearing is not None: bearing = round(bearing, 6)

        return bearing

class FCCollection:
    """Wrapper for the basic ArcGIS API operations."""
    featLayer = None

    members = []
    elementsModel = None

    selectAttrList = None
    editAttrList = None
    editableProperties = None

    aprxLayerName = None
    descr = None

    def __init__(self):
        self.members = []

    def readSelected(self, required_min=None, required_max=None):
        """
        Reads the selected features from the corresponding layer in the aprx project.
        Creates and stores the objects in the members list.

        :param required_min: required minimum (optional, integer)
        :param required_max: permitted maximum (optional, integer)
        """

        if self.descr is None: self.descr = arcpy.Describe(self.aprxLayerName)

        fid_set = self.descr.FIDset
        if fid_set != "":
            fid_list = fid_set.split("; ")
        else:
            fid_list = []

        nb_selected = len(fid_list)

        # validate the number of the selected features
        if required_min is not None and required_max is not None \
                and required_min == required_max and nb_selected != required_min:
            arcpy.AddError(f"Please select {required_min} features from \"{self.aprxLayerName}\" layer.")
            exit()
        elif required_min is not None and nb_selected < required_min:
            arcpy.AddError(f"Please select at least {required_min} features from \"{self.aprxLayerName}\" layer.")
            exit()
        elif required_max is not None and nb_selected > required_max:
            arcpy.AddError(f"Please select {required_min} or less features from \"{self.aprxLayerName}\" layer.")
            exit()

        if not fid_list: return

        # read the corresponding features
        object_ids = fid_set.replace(";", ",")
        self.readFromQuery(object_ids=object_ids)

    def readFromQuery(self, where_clause=None, object_ids=None):
        """
        Reads features from Query statement. Based on See ArcGIS REST APIs Query (Feature Service/Layer).
        Creates and stores the objects in the members list.

        :param where_clause: A WHERE clause for the query filter. Example: "OBJECTID in (29, 30)"
        :param object_ids: The object IDs of this layer or table to be queried. Example: "29, 30"
        """

        # verify the privileges
        query_capability = 'Query' in self.featLayer.properties.capabilities
        if not query_capability:
            arcpy.AddWarning(
                f'You do not have enough privileges to read from the layer "{self.featLayer.properties.name}"')

        # verify the parameters
        if where_clause is None and object_ids is None: return

        # query feature layer
        feature_set = self.featLayer.query(
            where=where_clause,
            objectIds=object_ids,
            out_fields=','.join(self.selectAttrList),
            return_geometry=True)

        # populate members list
        self.members = []
        for feat in feature_set.features:
            feat_obj = self.elementsModel(feat)
            self.members.append(feat_obj)

    def save(self, update_geom=True):
        """Saves the members of the collection in the corresponding feature layer.
        Based on arcgis.features module."""

        # verify the privileges
        query_capability = 'Update' in self.featLayer.properties.capabilities
        if not query_capability:
            arcpy.AddWarning(f'You do not have enough privileges to update the layer "{self.featLayer.properties.name}"')
            return

        for member in self.members:
            # get the properties of the member in a list

            values = [value for attr, value in member.__dict__.items() if attr in self.editableProperties]

            # actualiser les attributs de la Feature
            for i in range(0, len(self.editAttrList)):
                member.feat.set_value(field_name=self.editAttrList[i], value=values[i])

            # actualiser la géométrie
            if update_geom and hasattr(member, 'geom') and member.geom is not None:
                member.feat.geometry = member.geom

        try:
            arcpy.AddMessage([x.feat for x in self.members])
            result = self.featLayer.edit_features(updates=[x.feat for x in self.members])
            return result['updateResults']

        except Exception as e: arcpy.AddError(f"Write Failed: {e}")

    def selectMembers(self):

        if not self.members: return

        jointeur = ', '.join([str(x.id) for x in self.members])
        where_clause = f'OBJECTID in ({jointeur})'

        result = arcpy.SelectLayerByAttribute_management(self.aprxLayerName, "NEW_SELECTION", where_clause)
        while result.status < 4: sleep(0.2)


#!/usr/bin/env python2.7

from __future__ import print_function, division

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as clr
from matplotlib import cm
import dtk
import h5py
import sys
import os
import time
import healpy as hp
import astropy

from astropy.cosmology import WMAP7 as cosmo

npix = hp.nside2npix(32)

sky_frac = 131/npix
print("sky_frac: ", sky_frac)

print("here it is")
data_real = np.genfromtxt("count_real.txt", delimiter=' ')
data_uf = np.genfromtxt("count_uf.txt", delimiter=' ')
print(data_real.astype(np.int))

data_real_cumsum = np.cumsum(data_real[::-1,1], axis=0)
data_uf_cumsum = np.cumsum(data_uf[::-1,1], axis=0)

stepz = dtk.StepZ(sim_name='AlphaQ')
print(data_real[:,0])

redshifts = stepz.get_z(data_real[::-1,0])
comoving_vol = cosmo.comoving_volume(redshifts)*(0.7**3)

plt.figure()
plt.plot(redshifts, comoving_vol*sky_frac)
plt.xlabel('redshift')
plt.ylabel('volume [(h^-1 Mpc)^3]')
plt.figure()
plt.plot(redshifts, data_real_cumsum)
plt.plot(redshifts, data_uf_cumsum)
plt.xlabel('redshift')
plt.ylabel('accumlated normal galaxy count')

plt.figure()
plt.plot(redshifts, data_real_cumsum/(comoving_vol*sky_frac), label='normal')
plt.plot(redshifts, data_uf_cumsum/(comoving_vol*sky_frac), label = 'uf')
plt.xlabel('redshift')
plt.ylabel('galaxies/(h^-1 Mpc)^3')
plt.legend(loc='best')

plt.figure()
plt.plot(redshifts, data_real_cumsum/(comoving_vol*sky_frac), label='normal')
plt.plot(redshifts, data_uf_cumsum/(comoving_vol*sky_frac), label='uf')
plt.xlabel('redshift')
plt.ylabel('galaxies/(h^-1 Mpc)^3')
plt.yscale('log')
plt.legend(loc='best')

print(comoving_vol[-1]*sky_frac)
print(comoving_vol*sky_frac)

print(data_real_cumsum/(comoving_vol*sky_frac))

plt.figure()
plt.plot(redshifts[::-1], data_uf[:,1]/data_real[:,1])
plt.xlabel('redshift')
plt.ylabel('ultra faints galaxys/normal galaxies')
plt.show()


import tensorflow as tf
import numpy as np
from ops import *
from tensorflow.examples.tutorials.mnist import input_data

# 如果当前文件夹中没有mnist数据，则会自动下载
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)  # 使用独热编码，也就是标签是一个二维逻辑矩阵
train = mnist.train  # train有两个属性：images: 55000*784 和 labels： 55000*10

global_step = tf.Variable(0, name='global_step', trainable=False)
y = tf.placeholder(tf.float32, [BATCH_SIZE, 10], name='y')
images = tf.placeholder(tf.float32, [BATCH_SIZE, 28, 28, 1], name='real_images')
z = tf.placeholder(tf.float32, [BATCH_SIZE, 100], name='z')

# G是生成的假图片
with tf.variable_scope(tf.get_variable_scope()) as scope:
    G = generator(z, y)
    D, D_logits = discriminator(images, y)  # D、D_logits都是 BATCH_SIZE*1的
    D_, D_logits_ = discriminator(G, y, reuse=True)
    samples = sampler(z, y)

# 固定使用train.labels的前BATCH_SIZE个作为生成图片的标签，可以指定生成图片的数字
sample_labels = mnist.train.labels[0:BATCH_SIZE]

# 损失计算
d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=D_logits, labels=tf.ones_like(D)))
d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=D_logits_, labels=tf.zeros_like(D_)))
d_loss = d_loss_real + d_loss_fake
g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=D_logits_, labels=tf.ones_like(D_)))

# 生成器和判别器要更新的变量，用于 tf.train.Optimizer 的 var_list
t_vars = tf.trainable_variables()
d_vars = [var for var in t_vars if 'd_' in var.name]
g_vars = [var for var in t_vars if 'g_' in var.name]

# 由于使用了tf.layers.batch_normalization，需要添加下面的两行代码
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
    d_optim = tf.train.AdamOptimizer(0.0002, beta1=0.5).minimize(d_loss, var_list=d_vars, global_step=global_step)
    g_optim = tf.train.AdamOptimizer(0.0002, beta1=0.5).minimize(g_loss, var_list=g_vars, global_step=global_step)

with tf.Session() as sess:
    saver = tf.train.Saver()
    sess.run(tf.global_variables_initializer())
    for epoch in range(25):
        for i in range(int(55000/BATCH_SIZE)):
            batch = mnist.train.next_batch(BATCH_SIZE)
            batch_images = np.array(batch[0]).reshape((-1, 28, 28, 1))
            batch_labels = batch[1]
            batch_z = np.random.uniform(-1, 1, size=(BATCH_SIZE, 100))
            sess.run([d_optim], feed_dict={images: batch_images, z: batch_z, y: batch_labels})
            sess.run([g_optim], feed_dict={images: batch_images, z: batch_z, y: batch_labels})
            sess.run([g_optim], feed_dict={images: batch_images, z: batch_z, y: batch_labels})
            if i % 100 == 0:
                errD = d_loss.eval(feed_dict={images: batch_images, y: batch_labels, z: batch_z})
                errG = g_loss.eval({z: batch_z, y: batch_labels})
                print("epoch:[%d], i:[%d]  d_loss: %.8f, g_loss: %.8f" % (epoch, i, errD, errG))
            # 在训练过程中得到生成器生成的假的图片并保存
            if i % 100 == 1:
                sample = sess.run(samples, feed_dict={z: batch_z, y: sample_labels})
                samples_path = './pics/'
                save_images(sample, [8, 8], samples_path + 'epoch_%d_i_%d.png' % (epoch, i))
                print('save image')
            # 定期保存模型
            # if i == (int(55000/BATCH_SIZE)-1):
            #     checkpoint_path = os.path.join('./check_point/DCGAN_model.ckpt')
            #     saver.save(sess, checkpoint_path, global_step=i+1)
            #     print('save check_point')


import pygame
import sys
from pyanimation import Animation

pygame.init()

SCREEN_WIDTH, SCREEN_HEIGHT = 640, 480
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT), 0, 32)
surface = pygame.Surface(screen.get_size())
surface = surface.convert()
surface.fill((255,255,255))
clock = pygame.time.Clock()

link = Animation("images/link.png")
link.create_animation(0, 525, 120, 132, "front", duration=90, rows=1)
link.create_animation(0, 782, 120, 132, "back", duration=90, rows=1)

screen.blit(surface, (0, 0))

if __name__ == '__main__':
    while True:

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP:
                    link.run("back")
                if event.key == pygame.K_DOWN:
                    link.run("front")

        surface.fill((255,255,255))
        screen.blit(surface, (0,0))
        clock.tick(60)

        screen.blit(link.update_surface(), (link.x, link.y))
        pygame.display.flip()
        pygame.display.update()


import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import os
import time

class App(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("Vinke")

        self.minsize(1250, 900)

        self.option = tk.StringVar(self)
        self.options = ["Grid", "Columns", "Rows", "Noise", "Noise2", "Noise3", "Splice", "Weird Repeat", "Squares"]
        self.option.set(self.options[0])
        self.value = tk.IntVar(self)
        self.values = list(range(2,21))
        self.value.set(self.values[3])

        select_file_button1 = tk.Button(self, text="Select Image #1", command=self.select_file)
        select_file_button1.grid(row=0,column=0, sticky=tk.N,padx=10, pady=10)

        select_file_button2 = tk.Button(self, text="Select Image #2", command=self.select_file2)
        select_file_button2.grid(row=0,column=1, sticky=tk.N,padx=10, pady=10)

        option_menu = tk.OptionMenu(self, self.option, *self.options)
        option_menu.grid(row=0,column=2, sticky=tk.N,padx=10, pady=10)

        value_menu = tk.OptionMenu(self, self.value, *self.values)
        value_menu.grid(row=0,column=3, sticky=tk.N,padx=10, pady=10)

        generate_button = tk.Button(self, text="Generate", command=self.generate)
        generate_button.grid(row=0,column=4, sticky=tk.N,padx=10, pady=10)

        save_button = tk.Button(self, text="Save", command=self.save)
        save_button.grid(row=0,column=5, sticky=tk.N,padx=10, pady=10)

        # Create the canvas for displaying the image
        self.canvas1 = tk.Canvas(self, width=400, height=400, borderwidth=0, relief=tk.FLAT)
        self.canvas1.grid(row=1,column=0, sticky=tk.N,padx=10, pady=10)
        self.canvas2 = tk.Canvas(self, width=400, height=400, borderwidth=0, relief=tk.FLAT)
        self.canvas2.grid(row=2,column=0, sticky=tk.N,padx=10, pady=10)
        self.canvas_result  = tk.Canvas(self, width=800, height=800)
        self.canvas_result.grid(row=1,column=1, sticky=tk.N,padx=10, pady=10, rowspan=2, columnspan=5)

        self.image1 = None
        self.image2 = None
        self.image2_resized = None

        self.image_result = None

        dir_path = os.path.dirname(os.path.realpath(__file__))
        self.save_folder = os.fsencode(dir_path + '/saves/')
        self.examples_folder = os.fsencode(dir_path + '/examples/')

    def startup(self):
        dir_path = os.path.dirname(os.path.realpath(__file__))
        image1 = Image.open(os.fsencode(dir_path + '/examples/example8.jpg'))
        image2 = Image.open(os.fsencode(dir_path + '/examples/example9.jpg'))

        self.image1 = image1
        self.image2 = image2
        self.image2_resized = image2

        image_can1 = self.resize_image_for_canvas(self.image1, self.canvas1)
        image_can2 = self.resize_image_for_canvas(self.image2_resized, self.canvas2)

        self.display_image(image_can1, self.canvas1)
        self.display_image(image_can2, self.canvas2)

        self.generate()

    def select_file(self):
        # Open a file dialog to select an image file
        file_path = filedialog.askopenfilename(filetypes=[("Image files", "*.jpg;*.png")])

        # Check if a file was selected
        if file_path:
            # Load the image and display it on the canvas
            image = Image.open(file_path)
            self.image1 = image

            if self.image2 is not None:
                self.image2_resized = self.image2.resize(size=(self.image1.width, self.image1.height), resample=Image.LANCZOS)
                image_can2 = self.resize_image_for_canvas(self.image2_resized, self.canvas2)
                self.display_image(image_can2, self.canvas2)

            image = self.resize_image_for_canvas(image, self.canvas1)
            self.display_image(image, self.canvas1)

    def select_file2(self):
        # Open a file dialog to select an image file
        if self.image1 is None:
            messagebox.showerror("Error", "Select first image!")
            return

        file_path = filedialog.askopenfilename(filetypes=[("Image files", "*.jpg;*.png")])

        # Check if a file was selected
        if file_path:
            # Load the image and display it on the canvas
            image = Image.open(file_path)
            self.image2 = image

            image = image.resize(size=(self.image1.width, self.image1.height), resample=Image.LANCZOS)
            self.image2_resized = image

            image = self.resize_image_for_canvas(image, self.canvas2)

            self.display_image(image, self.canvas2)

    def generate(self):
        image1 = self.image1
        image2 = self.image2_resized

        if image1 is None or image2 is None:
            messagebox.showerror("Error", "Select both images!")
            return

        option = self.option.get()
        value = self.value.get()

        width = image1.width
        height = image1.height

        column_size = int(width/value)
        row_size = int(height/value)

        pixels1 = list(image1.getdata())
        pixels2 = list(image2.getdata())

        image = Image.new(image1.mode, image1.size)
        out_pixels = []

        if option == "Noise":
            for i in range(value):
                for j in range(value):
                    for x in range(column_size):
                        for y in range(row_size):
                            index = (j*row_size+y)*width + (i*column_size+x)
                            if (i+j) % 2 == 0:
                                out_pixels.append(pixels1[index])
                            else:
                                out_pixels.append(pixels2[index])
                            if i == value-1 and x == column_size-1:
                                for i in range(width-(i*column_size+x)-1):
                                    index += 1
                                    if (i+j) % 2 == 0:
                                        out_pixels.append(pixels1[index])
                                    else:
                                        out_pixels.append(pixels2[index])
        elif option == "Weird Repeat":
            for j in range(value):
                for i in range(value):
                    for y in range(row_size):
                        for x in range(column_size):
                            index = (j*row_size+y)*width + (i*column_size+x)
                            if (i+j) % 2 == 0:
                                out_pixels.append(pixels1[index])
                            else:
                                out_pixels.append(pixels2[index])
                            if i == value-1 and x == column_size-1:
                                for i in range(width-(i*column_size+x)-1):
                                    index += 1
                                    if (i+j) % 2 == 0:
                                        out_pixels.append(pixels1[index])
                                    else:
                                        out_pixels.append(pixels2[index])
        elif option == "Noise2":
            for x in range(column_size):
                for j in range(value):
                    for i in range(value):
                        for y in range(row_size):
                            index = (j*row_size+y)*width + (i*column_size+x)
                            if (i+j) % 2 == 0:
                                out_pixels.append(pixels1[index])
                            else:
                                out_pixels.append(pixels2[index])
                            if i == value-1 and x == column_size-1:
                                for i in range(width-(i*column_size+x)-1):
                                    index += 1
                                    if (i+j) % 2 == 0:
                                        out_pixels.append(pixels1[index])
                                    else:
                                        out_pixels.append(pixels2[index])
        elif option == "Noise3":
            for y in range(row_size):
                for x in range(column_size):
                    for i in range(value):
                        for j in range(value):
                            index = (j*row_size+y)*width + (i*column_size+x)
                            if (i+j) % 2 == 0:
                                out_pixels.append(pixels1[index])
                            else:
                                out_pixels.append(pixels2[index])
                            if i == value-1 and x == column_size-1:
                                for i in range(width-(i*column_size+x)-1):
                                    index += 1
                                    if (i+j) % 2 == 0:
                                        out_pixels.append(pixels1[index])
                                    else:
                                        out_pixels.append(pixels2[index])
        elif option == "Grid":
            for j in range(value):
                for y in range(row_size):
                    for i in range(value):
                        for x in range(column_size):
                            index = (j*row_size+y)*width + (i*column_size+x)
                            if (i+j) % 2 == 0:
                                out_pixels.append(pixels1[index])
                            else:
                                out_pixels.append(pixels2[index])
                            if i == value-1 and x == column_size-1:
                                for i in range(width-(i*column_size+x)-1):
                                    index += 1
                                    if (i+j) % 2 == 0:
                                        out_pixels.append(pixels1[index])
                                    else:
                                        out_pixels.append(pixels2[index])
        elif option == "Squares":
            for j in range(value):
                for y in range(row_size):
                    for i in range(value):
                        for x in range(column_size):
                            index = (j*row_size+y)*width + (i*column_size+x)
                            if i*j % 2 == 0:
                                out_pixels.append(pixels1[index])
                            else:
                                out_pixels.append(pixels2[index])
                            if i == value-1 and x == column_size-1:
                                for i in range(width-(i*column_size+x)-1):
                                    index += 1
                                    if i*j % 2 == 0:
                                        out_pixels.append(pixels1[index])
                                    else:
                                        out_pixels.append(pixels2[index])
        elif option == "Columns":
            for j in range(value):
                for y in range(row_size):
                    for i in range(value):
                        for x in range(column_size):
                            index = (j*row_size+y)*width + (i*column_size+x)
                            if i % 2 == 0:
                                out_pixels.append(pixels1[index])
                            else:
                                out_pixels.append(pixels2[index])
                            if i == value-1 and x == column_size-1:
                                for i in range(width-(i*column_size+x)-1):
                                    index += 1
                                    if i % 2 == 0:
                                        out_pixels.append(pixels1[index])
                                    else:
                                        out_pixels.append(pixels2[index])
        elif option == "Splice":
            for j in range(value):
                for y in range(row_size):
                    for i in range(value):
                        for x in range(column_size):
                            index = (j*row_size+y)*width + (i*column_size+x)
                            if x % 2 == 0:
                                out_pixels.append(pixels1[index])
                            else:
                                out_pixels.append(pixels2[index])
                            if i == value-1 and x == column_size-1:
                                for i in range(width-(i*column_size+x)-1):
                                    index += 1
                                    if x % 2 == 0:
                                        out_pixels.append(pixels1[index])
                                    else:
                                        out_pixels.append(pixels2[index])
        elif option == "Rows":
            for j in range(value):
                for y in range(row_size):
                    for i in range(value):
                        for x in range(column_size):
                            index = (j*row_size+y)*width + (i*column_size+x)
                            if j % 2 == 0:
                                out_pixels.append(pixels1[index])
                            else:
                                out_pixels.append(pixels2[index])
                            if i == value-1 and x == column_size-1:
                                for i in range(width-(i*column_size+x)-1):
                                    index += 1
                                    if j % 2 == 0:
                                        out_pixels.append(pixels1[index])
                                    else:
                                        out_pixels.append(pixels2[index])

        image.putdata(out_pixels)

        self.image_result = image

        image = self.resize_image_for_canvas(image, self.canvas_result)
        self.display_image(image, self.canvas_result)

    def save(self):
        timestr = time.strftime("%Y%m%d-%H%M%S")
        path = self.save_folder.decode("utf8") + timestr + ".jpg"

        self.image_result.save(path)

    def display_image(self, image, canvas):
        # Convert the image to a PhotoImage object
        image = ImageTk.PhotoImage(image)

        # Calculate the center coordinates of the canvas
        center_x = canvas.winfo_width() // 2
        center_y = canvas.winfo_height() // 2

        # Display the image on the canvas
        canvas.create_image(center_x, center_y, image=image, anchor=tk.CENTER)

        # Save a reference to the image to prevent it from being garbage collected
        canvas.image = image

    def get_image_on_canvas(self, canvas):
        # Get the list of items on the canvas
        items = canvas.find_all()

        # Check if there is an image on the canvas
        if not items:
            messagebox.showerror("Error", "No image to copy")
            return None

        # Get the image from the canvas
        image = canvas.image

        return image

    def resize_image_for_canvas(self, image, canvas):
        canvas_width = canvas.winfo_width()
        canvas_height = canvas.winfo_height()
        # Check if the image is larger than the canvas
        if image.width > canvas_width or image.height > canvas_height:
            # Calculate the new size of the image to fit the canvas
            width_ratio = canvas_width / image.width
            height_ratio = canvas_height / image.height
            resize_ratio = min(width_ratio, height_ratio)
            new_size = (int(image.width * resize_ratio), int(image.height * resize_ratio))
            # Resize the image
            image = image.resize(new_size, resample=Image.LANCZOS)

        return image

if __name__ == '__main__':
    app = App()
    app.after(200,app.startup)
    app.mainloop()


from sklearn.preprocessing import MinMaxScaler
from ydata_synthetic.synthesizers.timeseries import TimeGAN
from ydata_synthetic.preprocessing.timeseries.utils import real_data_loading
import numpy as np
import tensorflow as tf
import sys
import pickle
sys.path.insert(0, '../../data_process/')
from my_data_pre_intdash import DataPre

dataset_directory = '../../../datasets/'

def loadSynthData(model20, model40, model80, seq_len):
    synth_20 = TimeGAN.load(model20)
    synth_data_20 = synth_20.sample(seq_len)

    synth_40 = TimeGAN.load(model40)
    synth_data_40 = synth_40.sample(seq_len)

    synth_80 = TimeGAN.load(model80)
    synth_data_80 = synth_80.sample(seq_len)

    synth_data_20[:,:,14:18][synth_data_20[:,:,14:17] >= 0.5] = 1
    synth_data_20[:,:,14:18][synth_data_20[:,:,14:17] < 0.5] = 0

    synth_data_40[:,:,14:18][synth_data_40[:,:,14:17] >= 0.5] = 1
    synth_data_40[:,:,14:18][synth_data_40[:,:,14:17] < 0.5] = 0

    synth_data_80[:,:,14:18][synth_data_80[:,:,14:17] >= 0.5] = 1
    synth_data_80[:,:,14:18][synth_data_80[:,:,14:17] < 0.5] = 0

    return synth_data_20, synth_data_40, synth_data_80

def loadRealData(dsint20, dsint40, dsint80, dsdash20, dsdash40, dsdash80, num_cols, cat_cols, sample_size, randon, outliers):
    dp = DataPre()

    dp.loadDataSet(path20_int=dsint20, path40_int=dsint40, path80_int=dsint80, path20_dash=dsdash20, path40_dash=dsdash40, path80_dash=dsdash80)

    dp.preProcessData(num_cols, cat_cols=cat_cols, cond_col='q_size', random=randon)
    if outliers == False:
        dp.removeOutliers()

    real_data_20 = dp.processed_data.loc[dp.processed_data['q_size'] == 0].copy()
    real_data_20 = real_data_20[0:sample_size].copy()
    real_data_20 = real_data_20.values

    real_data_40 = dp.processed_data.loc[dp.processed_data['q_size'] == 1].copy()
    real_data_40 = real_data_40[0:sample_size].copy()
    real_data_40 = real_data_40.values

    real_data_80 = dp.processed_data.loc[dp.processed_data['q_size'] == 2].copy()
    real_data_80 = real_data_80[0:sample_size].copy()
    real_data_80 = real_data_80.values

    return real_data_20, real_data_40, real_data_80

def getStatistics(data):
    median = np.median(data)
    percentile_25 = np.percentile(data, 25)
    percentile_75 = np.percentile(data, 75)

    return [percentile_25, median, percentile_75]

def genStatisctics(real_20, synth_20, real_40, synth_40, real_80, synth_80, sample_size, num_cols):
    dict = {}

    for j, col in enumerate(num_cols):
        dict[col] = [getStatistics(real_20[:,j][:sample_size]),
                     getStatistics(synth_20[:,j][:sample_size]),
                     getStatistics(real_40[:,j][:sample_size]),
                     getStatistics(synth_40[:,j][:sample_size]),
                     getStatistics(real_80[:,j][:sample_size]),
                     getStatistics(synth_80[:,j][:sample_size])]

    return dict

def createDataSet(lines, seq_len, data):
    dataset = np.zeros(lines*seq_len*17).reshape(lines*seq_len,17)

    for i in range(0,lines):
        for j in range(0, seq_len):
            dataset[(i*seq_len) + j] = data[i][j][:]

    return dataset

def roundFeatures(dataset):
    for i in range(0, len(dataset)):
        dataset[i] = np.round(dataset[i])

def getMetrics(data):
    # [ (Upper_Quartile(Real) - Lower_Quartile(Real)) - (Upper_Quartile(Synth) - Lower_Quartile(Synth) ] + (Median(Real) - Median(Synth))

    metric20 = (abs( (data[0][2] - data[0][0]) - (data[1][2] - data[1][0])) +
               abs(  data[0][1] - data[1][1] ) )

    metric40 = (abs( (data[2][2] - data[2][0]) - (data[3][2] - data[3][0])) +
               abs(  data[2][1] - data[3][1] ) )

    metric80 = (abs( (data[4][2] - data[4][0]) - (data[5][2] - data[5][0])) +
               abs(  data[4][1] - data[5][1] ) )

    return metric20, metric40, metric80

def get_allfeatures_metrics(num_cols, arr, m, metrics):
    for j, col in enumerate(num_cols):
        arr[0][m][j], arr[1][m][j], arr[2][m][j] = getMetrics(metrics.get(col))

####################### GENERATION #######################

sample_size = 2400

num_cols = ['enq_qdepth1','deq_timedelta1', 'deq_qdepth1',
            ' enq_qdepth2', ' deq_timedelta2', ' deq_qdepth2',
            'enq_qdepth3', 'deq_timedelta3', 'deq_qdepth3',
            'Buffer', 'ReportedBitrate', 'FPS', 'CalcBitrate',
            'q_size']
cat_cols = ['Resolution']

real_20, real_40, real_80 = loadRealData(dsint20= dataset_directory + 'log_INT_20ms.csv',
                                        dsint40= dataset_directory + 'log_INT_40ms.csv',
                                        dsint80= dataset_directory + 'log_INT_80ms.csv',
                                        dsdash20= dataset_directory + 'dash_20ms.csv',
                                        dsdash40= dataset_directory + 'dash_40ms.csv',
                                        dsdash80= dataset_directory + 'dash_80ms.csv',
                                        num_cols=num_cols,
                                        cat_cols=cat_cols,
                                        sample_size=4000,
                                        randon=False,
                                        outliers=False)

scaler20 = MinMaxScaler().fit(real_20)
scaler40 = MinMaxScaler().fit(real_40)
scaler80 = MinMaxScaler().fit(real_80)

models = {}

size = 3600
data = np.zeros(3*9*len(num_cols)).reshape(3,9,len(num_cols))

try:
  # Specify an invalid GPU device
  with tf.device('/device:GPU:0'):
    for i in range(0,1):
        for j in range(0,3):
            for k in range(0,3):
                seq_len=(50*(i)+50)
                synth_20_norm, synth_40_norm, synth_80_norm = loadSynthData(model20= str('../saved_models/so20_seqlen_'+ str((50*(i) + 50)) + '_hidim_' + str(20*(j)+20) + '_batch_' +  str(28*(k) + 100) + '.pkl'),
                                                model40= str('../saved_models/so40_seqlen_'+ str((50*(i) + 50)) + '_hidim_' + str(20*(j)+20) + '_batch_' +  str(28*(k) + 100) + '.pkl'),
                                                model80= str('../saved_models/so80_seqlen_'+ str((50*(i) + 50)) + '_hidim_' + str(20*(j)+20) + '_batch_' +  str(28*(k) + 100) + '.pkl'),
                                                seq_len=int(size/seq_len))

                real_20_norm =  real_data_loading(real_20, seq_len=seq_len)
                real_40_norm =  real_data_loading(real_40, seq_len=seq_len)
                real_80_norm =  real_data_loading(real_80, seq_len=seq_len)

                real_20_norm = createDataSet(int(size/seq_len),seq_len, real_20_norm)
                real_40_norm = createDataSet(int(size/seq_len),seq_len, real_40_norm)
                real_80_norm = createDataSet(int(size/seq_len),seq_len, real_80_norm)

                synth_20_norm = createDataSet(int(size/seq_len),seq_len, synth_20_norm)
                synth_40_norm = createDataSet(int(size/seq_len),seq_len, synth_40_norm)
                synth_80_norm = createDataSet(int(size/seq_len),seq_len, synth_80_norm)

                data_synth_20 = scaler20.inverse_transform(synth_20_norm)
                data_synth_40 = scaler40.inverse_transform(synth_40_norm)
                data_synth_80 = scaler80.inverse_transform(synth_80_norm)

                roundFeatures(data_synth_20)
                roundFeatures(data_synth_40)
                roundFeatures(data_synth_80)

                models['so_seqlen_'+str((50*(i) + 50))+'_hidim_'+str(20*(j)+20)+'_batch_'+str(28*(k)+100)+'.pkl'] = [data_synth_20,
                                                                                data_synth_40,
                                                                                data_synth_80]
                metrics = genStatisctics(real_20_norm, synth_20_norm, real_40_norm, synth_40_norm, real_80_norm, synth_80_norm, sample_size, num_cols)

                print(metrics)
                get_allfeatures_metrics(num_cols, data, (i*9) + (j*3) + (k) , metrics)

    directory_path = '../saved_objects/'

    with open(directory_path + 'real9_50_3600_norm_27.pkl', 'wb') as file:
        pickle.dump([real_20, real_40, real_80], file)

    with open(directory_path + 'models9_50_3600_norm_27.pkl', 'wb') as file:
        pickle.dump(models, file)

    with open(directory_path + 'metrics9_50_3600_norm_27.pkl', 'wb') as file:
        pickle.dump(data, file)

except RuntimeError as e:
  print(e)