import numpy as np

def IQ_imbalance(g, phi):
    c = np.cos(phi)
    s = np.sin(phi)
    N = 1 / ((1 - g ** 2) * (2 * c ** 2 - 1))
    return [float(N * x) for x in [(1 - g) * c, (1 + g) * s,
                                   (1 - g) * s, (1 + g) * c]]

def catch_sech_shape(amp, total_time, kappa, lambd):
    if total_time < 16:
        pad_time = 16 - total_time
    else:
        pad_time = (4 - total_time % 4) % 4
    t = np.arange(total_time) - total_time // 2
    tau = kappa / 2. * t
    pulse = np.sqrt(lambd / 2. / (np.exp(lambd * tau) + 1. - lambd / 2.))
    pulse = pulse * (1 + lambd / 2. * np.tanh(lambd * tau / 2.))
    return list(amp * pulse) + [0.0] * pad_time

def sech_shape(amp, total_time, n_sigma, padding=True):
    pad_time = 0
    if padding:
        if total_time < 16:
            pad_time = 16 - total_time
        else:
            pad_time = (4 - total_time % 4) % 4
    t = np.array(range(-total_time // 2, total_time // 2, 1))
    sigma = total_time / n_sigma
    return list(amp / np.cosh(np.sqrt(np.pi / 2) * t / sigma)) + \
        [0.0] * pad_time

def ramsey_shape(amp1, amp2, time1, time2, time_wait, n_sigma):
    padding = (4 - (time1 + time2 + time_wait) % 4) % 4
    return sech_shape(amp1, time1, n_sigma, padding=False) + [0.0] * time_wait\
        + sech_shape(amp2, time2, n_sigma, padding=False) + [0.0] * padding

def ramsey_drag_shape(drag_amp1, drag_amp2, time1, time2, time_wait, n_sigma):

    padding = (4 - (time1 + time2 + time_wait) % 4) % 4
    return drag_sech(drag_amp1, time1, n_sigma=n_sigma) + [0.0] * time_wait\
        + drag_sech(drag_amp2, time2, n_sigma=n_sigma) + [0.0] * padding

def drag_sech(amp, total_time, n_sigma=4):
    t = np.array(range(-total_time // 2, total_time // 2, 1))
    sigma = total_time / n_sigma
    return list(-amp * np.tanh(np.sqrt(np.pi / 2) * t / sigma) /
                np.cosh(np.sqrt(np.pi / 2) * t / sigma))

def padded_time(t):
    return max(t + (4 - t % 4) % 4, 16)

parameters = {
    'sech_swap_amp': 0.1,
    'disp_amp': 0.0558,
    'phase_mod2_onoff': 1,
    'phase_mod4_onoff': 1,
    'correction_angle': -1.44 + 0.25,
    'buffer_sech_amp': 0.1,

    'delay_catch': 15,
}

optimized_parameters = {
    'offset_memory': 0.017,
    'offset_buffer_I': -0.0037,
    'offset_buffer_Q': -0.0016,
    'offset_pump_I': -0.0047,
    'offset_pump_Q': -0.0084,
    'offset_ampli': 0.,
    'offset_qubit_I': -0.0013,
    'offset_qubit_Q': -0.0038,
    'drag_amp': 0.05,

    'qubit_if': 200e6,
    'qubit_long_if':200e6,

    'wait_mod2': 140,
    'wait_mod4': 64,
    'buffer_I': 1,
    'unconditional_pi_2_length': 16,
    'unconditional_pi_2_amp': 0.2267,
    'unconditional_pi_length': 32,
    'unconditional_pi_amp': 0.2,

    'buffer_sech_length': 300,
    'theta': 0,

    'y_scaling': 1,

    'disp_length': 52,
    'sech_catch_amp': 0.165,
    'sech_catch_length': 300,

    'lambd': 0.532,
    'kappa': 0.02 * 2 * np.pi,
    'ki': 0.001 * 2 * np.pi,
    'meas_length': 252,

    'sech_swap_length': 240,

    'memory_if': 120e6,
    'pump_if': 170e6,
    'buffer_if': 50e6,

    'short_ro_length': 252,
    'short_ro_amp': 0.16,

    'clear_ro_amp_up': 0.16,
    'clear_ro_amp_steady': 0,
    'clear_ro_amp_down': -0.115,
    'clear_ro_amp_down_steady': 0.055,
    'clear_ro_length_up': 252,
    'clear_ro_length_steady': 0,
    'clear_ro_length_down': 100,
    'clear_ro_length_down_steady': 52,

    'ro_if': 51e6,
    'weights_path': 'weights.npy',

    'square_swap_amp': 0.01,
    'square_swap_length': 200,
    'buffer_arbitrary_amp': 0.1,
    'selective_pi_length': 560,
    'selective_pi_amp': 0.0119,
}

def get_parameters():
    return parameters

def get_config(params):
    params.update(optimized_parameters)

    offset_qubit_I = float(params['offset_qubit_I'])
    offset_qubit_Q = float(params['offset_qubit_Q'])
    offset_memory = float(params['offset_memory'])
    offset_buffer_I = float(params['offset_buffer_I'])
    offset_buffer_Q = float(params['offset_buffer_Q'])
    offset_pump_I = float(params['offset_pump_I'])
    offset_pump_Q = float(params['offset_pump_Q'])
    offset_ampli = float(params['offset_ampli'])

    drag_amp = float(params['drag_amp'])
    y_scaling = float(params['y_scaling'])
    theta = float(params['theta'])

    qubit_if = float(params['qubit_if'])
    qubit_long_if = float(params['qubit_long_if'])
    buffer_if = float(params['buffer_if'])
    pump_if = float(params['pump_if'])
    memory_if = float(params['memory_if'])

    ro_if = float(params['ro_if'])

    clear_ro_length_up = int(params['clear_ro_length_up'])
    clear_ro_amp_up = float(params['clear_ro_amp_up'])
    clear_ro_length_steady = int(params['clear_ro_length_steady'])
    clear_ro_amp_steady = float(params['clear_ro_amp_steady'])

    clear_ro_length_down = int(params['clear_ro_length_down'])
    clear_ro_amp_down = float(params['clear_ro_amp_down'])
    clear_ro_length_down_steady = int(params['clear_ro_length_down_steady'])
    clear_ro_amp_down_steady = float(params['clear_ro_amp_down_steady'])

    short_ro_length = int(params['short_ro_length'])
    short_ro_amp = float(params['short_ro_amp'])

    unconditional_pi_length = int(params['unconditional_pi_length'])
    unconditional_pi_amp = float(params['unconditional_pi_amp'])
    unconditional_pi_2_length = int(params['unconditional_pi_2_length'])
    unconditional_pi_2_amp = float(params['unconditional_pi_2_amp'])
    selective_pi_length = int(params['selective_pi_length'])
    selective_pi_amp = float(params['selective_pi_amp'])

    meas_length = int(params['meas_length'])

    buffer_I = float(params['buffer_I'])
    buffer_arbitrary_amp = float(params['buffer_arbitrary_amp'])

    wait_mod2 = int(params['wait_mod2'])
    wait_mod4 = int(params['wait_mod4'])
    phase_mod2_onoff = int(params['phase_mod2_onoff'])
    phase_mod4_onoff = int(params['phase_mod4_onoff'])

    correction_angle = float(params['correction_angle'])
    if params['weights_path'] != '':
        weights = np.load(params['weights_path'])
    else:
        weights = np.array([1.0] * (padded_time(meas_length) // 4))
    weights1 = (np.cos(correction_angle) * weights).tolist()
    weights2 = (np.sin(correction_angle) * weights).tolist()

    square_swap_amp = float(params['square_swap_amp'])
    square_swap_length = int(params['square_swap_length'])
    sech_catch_amp = float(params['sech_catch_amp'])
    sech_catch_length = int(params['sech_catch_length'])
    sech_swap_amp = float(params['sech_swap_amp'])
    sech_swap_length = int(params['sech_swap_length'])

    buffer_sech_amp = float(params['buffer_sech_amp'])
    buffer_sech_length = int(params['buffer_sech_length'])

    delay_catch = int(params['delay_catch'])
    padding_catch = (4 - delay_catch % 4) % 4

    lambd = float(params['lambd'])
    kappa = float(params['kappa'])
    ki = float(params['ki'])

    random_shape = np.load('arbitrary.npy')

    config1 = {
        'version': 1,

        'controllers': {
            'con1': {
                'type': 'opx1',
                'analog_outputs': {
                    1: {'offset': offset_qubit_I},
                    2: {'offset': offset_qubit_Q},
                    4: {'offset': offset_memory},
                    5: {'offset': offset_buffer_I},
                    6: {'offset': offset_buffer_Q},
                    7: {'offset': offset_pump_I},
                    8: {'offset': offset_pump_Q},
                },
                'digital_outputs': {
                    2: {}
                },
                'analog_inputs': {
                    1: {'offset': 0.1945},
                    2: {'offset': 0.196},
                }
            }
        },

        'elements': {
            'qubit_short': {
                'mixInputs': {
                    'I': ('con1', 1),
                    'Q': ('con1', 2),
                    'mixer': 'mixer_qubit_short',
                    'lo_frequency': 4.4792e9
                },
                'digitalInputs': {
                    'port': {
                        'port': ('con1', 2),
                        'buffer': 0,
                        'delay': 130
                    }
                },
                'intermediate_frequency': qubit_if,
                'operations': {
                    'unconditional_pi': 'unconditional_pi_pulse',
                    'mod_2': 'mod_2_pulse',
                    'mod_4': 'mod_4_pulse',
                    'mod_4_y': 'mod_4_y_pulse',
                }
            },
            'qubit_long': {
                'mixInputs': {
                    'I': ('con1', 1),
                    'Q': ('con1', 2),
                    'mixer': 'mixer_qubit_long',
                    'lo_frequency': 4.4792e9
                },
                'digitalInputs': {
                    'port': {
                        'port': ('con1', 2),
                        'buffer': 0,
                        'delay': 130
                    }
                },
                'intermediate_frequency': qubit_long_if,
                'operations': {
                    'selective_pi': 'selective_pi_pulse',
                }
            },
            'readout': {
                'singleInput': {
                    'port': ('con1', 3)
                },
                'intermediate_frequency': ro_if,
                'operations': {
                    'meas_clear': 'meas_clear_pulse',
                    'clear_ro': 'clear_ro_pulse',
                },
                'outputs': {
                    'out1': ('con1', 2)
                },
                'digitalInputs': {
                    'port': {
                        'port': ('con1', 2),
                        'buffer': 0,
                        'delay': 130
                    }
                },
                'time_of_flight': 56,
                'smearing': 0,
            },
            'pump': {
                'mixInputs': {
                    'I': ('con1', 7),
                    'Q': ('con1', 8),
                    'mixer': 'mixer_pump',
                    'lo_frequency': 10.222e9 - 3.627e9 + 170e6
                },
                'digitalInputs': {
                    'port': {
                        'port': ('con1', 2),
                        'buffer': 0,
                        'delay': 130
                    }
                },
                'intermediate_frequency': pump_if,
                'operations': {
                    'sech_catch': 'sech_catch_pulse',
                    'sech_swap': 'sech_swap_pulse',
                    'square_swap': 'square_swap_pulse',
                }
            },
            'buffer': {
                'mixInputs': {
                    'I': ('con1', 5),
                    'Q': ('con1', 6),
                    'mixer': 'mixer_buffer',
                    'lo_frequency': 10.222e9 + 50e6
                },
                'digitalInputs': {
                    'port': {
                        'port': ('con1', 2),
                        'buffer': 0,
                        'delay': 130
                    }
                },
                'intermediate_frequency': buffer_if,
                'operations': {
                    'sech_kick': 'buffer_sech_pulse',
                    'arbitrary':'arbitrary_pulse'
                },
                'outputs': {
                    'out1': ('con1', 1)
                },
                'time_of_flight': 48,
                'smearing': 0
            },
        },
        'pulses': {
            'mod_2_pulse': {
                'operation': 'control',
                'length': padded_time(wait_mod2 + 2 * unconditional_pi_2_length),
                'waveforms': {
                    'I': 'mod2_I_wf',
                    'Q': 'mod2_Q_wf'
                },
                'digital_marker': 'trigger'
            },
            'mod_4_pulse': {
                'operation': 'control',
                'length': padded_time(wait_mod4 + 2 * unconditional_pi_2_length),
                'waveforms': {
                    'I': 'mod4_I_wf',
                    'Q': 'mod4_Q_wf'
                },
                'digital_marker': 'trigger'
            },
            'mod_4_y_pulse': {
                'operation': 'control',
                'length': padded_time(wait_mod4 + 2 * unconditional_pi_2_length),
                'waveforms': {
                    'I': 'mod4_y_I_wf',
                    'Q': 'mod4_y_Q_wf'
                },
                'digital_marker': 'trigger'
            },
            'unconditional_pi_pulse': {
                'operation': 'control',
                'length': padded_time(unconditional_pi_length),
                'waveforms': {
                    'I': 'unconditional_pi_wf',
                    'Q': 'zero_wf'
                },
                'digital_marker': 'trigger'
            },
            'clear_ro_pulse': {
                'operation': 'measurement',
                'length': padded_time(clear_ro_length_up + clear_ro_length_steady),
                'waveforms': {
                    'single': 'clear_ro_wf'
                },
                'integration_weights': {
                },
                'digital_marker': 'ON'
            },
            'meas_clear_pulse': {
                'operation': 'measurement',
                'length': padded_time(meas_length),
                'waveforms': {
                    'single': 'meas_clear_wf'
                },
                'integration_weights': {
                    'optiW1': 'opti_clear_W1',
                    'optiW2': 'opti_clear_W2',
                },
                'digital_marker': 'ON'
            },
            'sech_catch_pulse': {
                'operation': 'control',
                'length': padded_time(sech_catch_length),
                'waveforms': {
                    'I': 'sech_catch_wf',
                    'Q': 'zero_wf'
                },
                'digital_marker': 'trigger'
            },
            'sech_swap_pulse': {
                'operation': 'control',
                'length': padded_time(sech_swap_length),
                'waveforms': {
                    'I': 'sech_swap_wf',
                    'Q': 'zero_wf'
                },
                'digital_marker': 'trigger'
            },
            'buffer_sech_pulse': {
                'operation': 'control',
                'length': padded_time(buffer_sech_length) + delay_catch + padding_catch,
                'waveforms': {
                    'I': 'buffer_sech_I_wf',
                    'Q': 'buffer_sech_Q_wf'
                },
                'digital_marker': 'trigger'
            },
            'selective_pi_pulse': {
                'operation': 'control',
                'length': padded_time(selective_pi_length),
                'waveforms': {
                    'I': 'selective_pi_wf',
                    'Q': 'zero_wf'
                },
                'digital_marker': 'trigger'
            },
            'square_swap_pulse': {
                'operation': 'control',
                'length': padded_time(square_swap_length),
                'waveforms': {
                    'I': 'square_swap_wf',
                    'Q': 'zero_wf'
                },
                'digital_marker': 'trigger'
            },
            'arbitrary_pulse': {
                'operation': 'control',
                'length': padded_time(len(random_shape)),
                'waveforms': {
                    'I': 'arbitrary_wf',
                    'Q': 'zero_wf'
                },
                'digital_marker': 'ON'
            },
        },
        'waveforms': {
            'short_ro_wf': {
                'type': 'constant',
                'sample': short_ro_amp
            },
            'clear_ro_wf': {
                'type': 'arbitrary',
                'samples': [clear_ro_amp_up] * clear_ro_length_up\
                + [clear_ro_amp_steady] * clear_ro_length_steady\
            },
            'meas_clear_wf': {
                'type': 'arbitrary',
                'samples': [clear_ro_amp_down] * clear_ro_length_down\
                + [clear_ro_amp_down_steady] * clear_ro_length_down_steady\
                + [0] * (padded_time(meas_length) - clear_ro_length_down - clear_ro_length_down_steady)
            },
            'zero_wf': {
                'type': 'constant',
                'sample': 0.0
            },
            'mod2_I_wf': {
                'type': 'arbitrary',
                'samples': ramsey_shape(unconditional_pi_2_amp,
                                        (-1.)**int(phase_mod2_onoff) * unconditional_pi_2_amp,
                                        unconditional_pi_2_length, unconditional_pi_2_length, wait_mod2, 4)
            },
            'mod2_Q_wf': {
                'type': 'arbitrary',
                'samples': ramsey_drag_shape(drag_amp,
                                             (-1.)**int(phase_mod2_onoff) * drag_amp,
                                             unconditional_pi_2_length, unconditional_pi_2_length, wait_mod2, 4)
                #
            },
            'mod4_I_wf': {
                'type': 'arbitrary',
                'samples': list(np.array(ramsey_shape(unconditional_pi_2_amp,
                                                      (-1)**phase_mod4_onoff * unconditional_pi_2_amp * (1 - phase_mod2_onoff),
                                                      unconditional_pi_2_length, unconditional_pi_2_length, wait_mod4, 4))\
                                + np.array(ramsey_drag_shape(0,
                                                             -(-1)**phase_mod4_onoff * drag_amp * (phase_mod2_onoff),
                                                             unconditional_pi_2_length, unconditional_pi_2_length, wait_mod4, 4)))
            },
            'mod4_Q_wf': {
                'type': 'arbitrary',
                'samples': list(np.array(ramsey_shape(0,
                                                      (-1)**phase_mod4_onoff * unconditional_pi_2_amp * phase_mod2_onoff,
                                                      unconditional_pi_2_length, unconditional_pi_2_length, wait_mod4, 4))\
                                + np.array(ramsey_drag_shape(drag_amp,
                                                             (-1)**phase_mod4_onoff * drag_amp * (1 - phase_mod2_onoff),
                                                             unconditional_pi_2_length, unconditional_pi_2_length, wait_mod4, 4)))
            },

            # mod4_y can be on x depending on phase_mod2_onoff...
            'mod4_y_I_wf': {
                'type': 'arbitrary',
                'samples': list(np.array(ramsey_shape(unconditional_pi_2_amp,
                                                      (-1)**phase_mod4_onoff * unconditional_pi_2_amp * phase_mod2_onoff,
                                                      unconditional_pi_2_length, unconditional_pi_2_length, wait_mod4, 4))\
                                + np.array(ramsey_drag_shape(-drag_amp,
                                                             -(-1)**phase_mod4_onoff * drag_amp * (1 - phase_mod2_onoff),
                                                             unconditional_pi_2_length, unconditional_pi_2_length, wait_mod4, 4)))
            },
            'mod4_y_Q_wf': {
                'type': 'arbitrary',
                'samples': list(np.array(ramsey_shape(0, (-1)**phase_mod4_onoff * unconditional_pi_2_amp * (1 - phase_mod2_onoff),
                                                      unconditional_pi_2_length, unconditional_pi_2_length, wait_mod4, 4))\
                                + np.array(ramsey_drag_shape(drag_amp,
                                                             (-1)**phase_mod4_onoff * drag_amp * (phase_mod2_onoff),
                                                             unconditional_pi_2_length, unconditional_pi_2_length, wait_mod4, 4)))
            },
            'unconditional_pi_wf': {
                'type': 'arbitrary',
                'samples': sech_shape(unconditional_pi_amp, unconditional_pi_length, 4)
            },
            'sech_catch_wf': {
                'type': 'arbitrary',
                'samples': catch_sech_shape(sech_catch_amp, sech_catch_length, kappa, lambd)
            },
            'sech_swap_wf': {
                'type': 'arbitrary',
                'samples': sech_shape(sech_swap_amp, sech_swap_length, 4)
            },
            'buffer_sech_wf': {
                'type': 'arbitrary',
                'samples': [0.0] * delay_catch + sech_shape(buffer_sech_amp, buffer_sech_length, 4) \
                + [0.0] * padding_catch
            },
            'buffer_sech_I_wf': {
                'type': 'arbitrary',
                'samples': list((buffer_I) * np.array([0.0] * delay_catch + sech_shape(buffer_sech_amp, buffer_sech_length, 4) \
                                                      + [0.0] * padding_catch))
            },
            'buffer_sech_Q_wf': {
                'type': 'arbitrary',
                'samples': list((1 - buffer_I) * np.array([0.0] * delay_catch + sech_shape(buffer_sech_amp, buffer_sech_length, 4) \
                                                          + [0.0] * padding_catch))
            },
            'square_swap_wf': {
                'type': 'constant',
                'sample': square_swap_amp
            },
            'arbitrary_wf': {
                'type': 'arbitrary',
                'samples': list(np.array(random_shape)*buffer_arbitrary_amp) + [0.0]*(padded_time(len(random_shape))-len(random_shape))
            },
            'selective_pi_wf': {
                'type': 'arbitrary',
                'samples': sech_shape(selective_pi_amp, selective_pi_length, 4)
            },
        },
        'digital_waveforms': {
            'trigger': {
                'samples': [(1, 20), (0, 0)]
            },
            'ON': {
                'samples': [(1, 0)]
            }
        },
        'integration_weights': {
            'optiW1': {
                'cosine': weights1,
                'sine': [-i for i in weights2]
            },
            'optiW2': {
                'cosine': weights2,
                'sine': weights1
            },
        },

        'mixers': {
            'mixer_qubit_short': [{
                'intermediate_frequency': qubit_if,
                'correction': IQ_imbalance(0.255, np.pi * 0.011),
                'lo_frequency': 4.4792e9,
            }],
            'mixer_pump': [{
                'intermediate_frequency': pump_if,
                'correction': IQ_imbalance(0.005, np.pi * 0.014),
                'lo_frequency': 10.222e9 - 3.627e9 + 170e6,
            }],
            'mixer_buffer': [{
                'intermediate_frequency': buffer_if,
                'correction': IQ_imbalance(-0.001, -np.pi * 0.054),
                'lo_frequency': 10.222e9 + 50e6,
            }],
            'mixer_qubit_long': [{
                'intermediate_frequency': qubit_long_if,
                'correction': IQ_imbalance(0.088, -np.pi * 0.04),
                'lo_frequency': 4.4792e9,
            }],
        }
    }
    return config1


from contextlib import contextmanager
from django.core.management.base import BaseCommand
from os.path import basename
from time import time
import boto3
import crowdcount.models.paths as ccp
import requests

class Command(BaseCommand):
    def handle(self, *args, **kwargs):
        now = int(time())
        source = "https://cdn.shakeshack.com/camera.jpg?{}".format(now)
        destination = ccp.output("shakeshack-{}.jpg".format(now))

        with _fetch_image(source) as data:
            if len(data) > 0:
                _upload(data, "shakeshack/{}".format(basename(destination)))

@contextmanager
def _fetch_image(source):
    response = requests.get(source, stream=True)

    # Throw an error for bad status codes
    response.raise_for_status()
    yield response.content

def _upload(data, key):
    print("Uploading {}...".format(key))
    s3 = boto3.resource('s3')
    s3.Bucket('dev.counting-company.com').put_object(Key=key, Body=data)


import numpy as np
from scipy import optimize
from scipy.interpolate import Rbf, InterpolatedUnivariateSpline
import matplotlib.pyplot as plt
import time

class cftiming:
    def __init__(self):
        self.gain = 0.5
        self.delay = 1000 # delay 1000 ps
        self.method = 'tomas'

    def normalizeY(self, points):
        scale = points[np.argmax(np.abs(points))]
        return np.multiply(points, 1.0/scale)

    def crossFinder(self, y):
        jmin = y.argmin()
        jmax = y.argmax()
        x = np.where(np.diff(np.sign(y)))[0]
        arr1 = np.where(x<jmax)
        if len(arr1[0]) == 0: return -2
        arr1 = arr1[0][-1]
        return int(x[arr1])

    def simulateCFT(self, xs, ys, dt):
        frac = self.gain
        f = InterpolatedUnivariateSpline(xs, ys)
        def simCFT( x):
            return f(x-self.delayInterval*dt)-frac*f(x)
        return simCFT

    def getCFTiming(self, t, y):
        y = self.normalizeY(y)
        dt = t[1]-t[0]
        self.delayInterval = int(self.delay/dt)
        ntrun =int(np.floor(2000/dt))
        ymax = np.amax(y)
        halfy = np.where(y > ymax*self.gain*0.9)[0]
        t0 = halfy[0]
        t1 = halfy[-1]+self.delayInterval
        f = self.simulateCFT(t, y, dt)
        yy = f(t[ntrun:-ntrun])
        yys = yy[t0-ntrun:t1-ntrun]
        if len(yys) == 0 : print(len(yy), [t0-ntrun,t1-ntrun])
        x0 = self.crossFinder(yys)
        if x0 < 0 : return -1
        inte = [t[t0+x0-1], t[t0+x0+1]]
        if self.method == 'tomas':
            try:
                cft = optimize.toms748(f, inte[0],inte[1])
            except:
                print(inte)
                raise ValueError('failed to get CFT')
            #cft = optimize.newton(lambda x : f(x), x0=(inte[0]+inte[1])/2)
        elif self.method == 'linear':
            x1 = inte[0]
            x2 = inte[1]
            k = (f(x2)-f(x1))/(x2-x1)
            cft = x1-f(x1)/k
        return cft

    def debug(self, t, y):
        y = self.normalizeY(y)
        dt = t[1]-t[0]
        ntrun =int(np.floor(2000/dt))
        ymax = np.amax(y)
        halfy = np.where(y > ymax*self.gain*0.9)[0]
        t0 = halfy[0]
        t1 = halfy[-1]+self.delayInterval
        f = self.simulateCFT(t, y, dt)
        yy = f(t[ntrun:-ntrun])
        yys = yy[t0-ntrun:t1-ntrun]
        plt.plot(t[t0-100:t1+100], f(t[t0-100:t1+100]))
        plt.plot(t[t0:t1], f(t[t0:t1]))
        plt.show()

    def getStrange(self, dt):
        print(np.argmax(np.abs(dt)), np.amax(np.abs(dt)))

def gaussian(x, mean, sig, amp):
    return amp*exp(-(x-mean)**2/sig**2)


from core.loader.loader import uLoad

import logging
log = logging.getLogger("msg")
log.setLevel(logging.DEBUG)

class MessageAction(uLoad):
    """'
    Just print ALL message form MBUS in console
    """

    def _activate(self):

        self.sub_h(topic="ALL", func="print_mbus")
        self.mbus.pub_h("module", "message")

    @staticmethod
    def print_mbus(*msg):
        print("MBUS: {}".format(msg))


def has_index(ctx):
    return ctx.have_indexables

def get_index(ctx):
    for s in ctx.syms.sorted_name:
        if not s[2]:
            continue # skip symbols that don't want to appear in the index
        if not ctx.is_destination(s[0].first):
            continue # skip if we've got nothing to link to
        ctx.link_add_referenced(s[0].first) # reference so the anchor is generated
        yield {
            'name': s[1],
            'address': s[0].first
        }


import pymysql.cursors
from PyQt5 import uic,QtWidgets
from time import sleep as delay

SqlConexao = pymysql.connect(
            host = 'localhost',
            user = 'root',
            password= '',
            database = 'estacionamento',
            cursorclass= pymysql.cursors.DictCursor)

# SELECT (((saida - entrada)/100)*6) FROM carros WHERE placa ='PNL1234'; PAGAMENTO VAI SER IMPLEMENTADO

def Consultar():
    ConsultaSQL = SqlConexao.cursor()
    Consulta = ('SELECT placa FROM carros WHERE placa = %s')
    placa = input('Informe a placa do veiculo: ')
    print("#" * 10)
    print('Consultando...')
    print("#" * 10)
    delay(1.3)
    ConsultaSQL.execute(Consulta,({placa}))
    ApresentaConsulta = ConsultaSQL.fetchall()
    for Placa in ApresentaConsulta:
        print("#" * 10)
        print(f'{Placa} - PLACA CADASTRADA')
        print("#" * 10)
        op = int(input('1 - Nova Consulta | 2 - Retornar ao Menu : '))
        if op == 1:
            Consultar()
        if op == 2:
            Return()
    delay(0.4)
    print("#" * 10)
    print(f'{placa} - NAO CADASTRADA')
    print("#" * 10)
    op = int(input('1 - Nova Consulta | 2 - Retornar ao Menu : '))
    if op == 1:
        Consultar()
    if op == 2:
        Return()
    ConsultaSQL.close()
    SqlConexao.close()

def Apagar():
    ApagaSql = SqlConexao.cursor()
    Apaga = ('DELETE FROM carros WHERE placa = %s')
    placa = input("Informe a placa do veiculo:")
    ApagaSql.execute(Apaga,({placa}))
    confirma = int(input(f'Os a serem removidos são {placa} . Deseja confirmar? (1 - SIM / 2 - NÃO) '))
    if confirma == 1:
        SqlConexao.commit()
    else:
        Apagar()
    ApagaSql.close()
    SqlConexao.close()

def Modificar():
    ModificaSql = SqlConexao.cursor()
    Modifica = ("UPDATE carros SET placa = %s , modelo = %s WHERE cod =%s")
    placa = input("Informe a placa do veiculo: ")
    modelo = input("Informe o modelo do veiculo: ")
    cod = int(input("Informe o cod de registro : "))
    ModificaSql.execute(Modifica, ({placa}, {modelo},{cod}))
    confirma = int(input(f'Os dados serão atualizados {placa} e {modelo}. Deseja confirmar? (1 - SIM / 2 - NÃO) '))
    if confirma == 1:
        SqlConexao.commit()
    else:
        Modificar()
    ModificaSql.close()
    SqlConexao.close()

def Todos():
    ConsultaSQL = SqlConexao.cursor()
    Listando = ('SELECT * FROM carros')
    print('Consultando...')
    delay(1.3)
    ConsultaSQL.execute(Listando)
    ApresentaConsulta = ConsultaSQL.fetchall()
    for Placa in ApresentaConsulta:
        print(Placa)
    ConsultaSQL.close()
    SqlConexao.close()

def RegistroOperador():
    TelaCadastro.show()
    placa = TelaCadastro.Placa_Line.text()
    HoraEntrada = TelaCadastro.Entrada_Line.text()
    OperadorSql = SqlConexao.cursor()
    OperadorEntrada = ("INSERT INTO carros(placa,entrada)"
                "VALUES (%s,%s)")
    info = (str(placa),str(HoraEntrada))
    OperadorSql.execute(OperadorEntrada,info)
#    confirma = int(input(f'Os dados cadastrados são {placa} e {HoraEntrada}. Deseja confirmar? (1 - SIM / 2 - NÃO)'))
#    if confirma == 1:
    confirmar = SqlConexao.commit()
    PRINCIPAL.TelaCadastro.Confirma_btt.clicked.connect(confirmar)
    print(f'Os dados cadastrados são {placa} e {HoraEntrada}')
#       print('CADASTRANDO...')
#        delay(0.3)
#       Return()
#    else:
#        RegistroOperador()
    OperadorSql.close()
    SqlConexao.close()

def SaidaOperador():
    Consultar()
    delay(0.9)
    SaidaSql = SqlConexao.cursor()
    Saida = ("INSERT INTO carros(saida)"
                "VALUES (%s)")
    HoraSaida = input("Informe o modelo do veiculo:")
    SaidaSql.execute(Saida, ({HoraSaida},))
    confirma = int(input(f'Os dados cadastrados são {HoraSaida} . Deseja confirmar? (1 - SIM / 2 - NÃO)'))
    if confirma == 1:
        SqlConexao.commit()
        Return()
    else:
        SaidaOperador()
    SaidaSql.close()
    SqlConexao.close()

def ConsultaOperador():
    ConsultaSQL = SqlConexao.cursor()
    Consulta = ('SELECT placa,entrada FROM carros WHERE placa = %s')
    placa = input('Informe a placa do veiculo: ')
    print("#" * 10)
    print('Consultando...')
    print("#" * 10)
    delay(1.3)
    ConsultaSQL.execute(Consulta, ({placa}))
    ApresentaConsulta = ConsultaSQL.fetchall()
    for Placa in ApresentaConsulta:
        print("#" * 10)
        print(f'{Placa} - PLACA CADASTRADA')
        print("#" * 10)
        op = int(input('1 - Nova Consulta | 2 - Retornar ao Menu : '))
        if op == 1:
            Consultar()
        if op == 2:
            Return()
    delay(0.4)
    print("#" * 10)
    print(f'{placa} - NAO CADASTRADA')
    print("#" * 10)
    op = int(input('1 - Nova Consulta | 2 - Retornar ao Menu : '))
    if op == 1:
        Consultar()
    if op == 2:
        Return()
    ConsultaSQL.close()
    SqlConexao.close()

def Return():
    import OPERADOR
    OPERADOR.Menu()


#Programa para encontrar o maior e o menor valor da uma sequência de números
largest = None
smallest = None
while True:
    num = input("Enter a number: ")
    if num == "done":
        break
    try:
        value = int(num)
        if largest is None:
            largest = value
        elif value > largest:
            largest = value
        elif smallest is None:
            smallest = value
        elif value < smallest:
            smallest = value
    except:
        print('Invalid input')
print("Maximum is", largest)
print("Minimum is", smallest)


"""Add Budget tables

Revision ID: 50bc2928c41c
Revises:
Create Date: 2022-07-15 02:33:54.192335

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '50bc2928c41c'
down_revision = None
branch_labels = None
depends_on = None

def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('budget_amount',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('month', sa.Integer(), nullable=True),
    sa.Column('year', sa.Integer(), nullable=True),
    sa.Column('amount', sa.Numeric(precision=15, scale=2), nullable=True),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('budget_category',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(length=128), nullable=False),
    sa.PrimaryKeyConstraint('id'),
    sa.UniqueConstraint('name')
    )
    op.create_table('budget_item',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('value', sa.Numeric(precision=15, scale=2), nullable=True),
    sa.PrimaryKeyConstraint('id')
    )
    # ### end Alembic commands ###

def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('budget_item')
    op.drop_table('budget_category')
    op.drop_table('budget_amount')
    # ### end Alembic commands ###


"""Slice a Volume with multiple planes
Make low values of the scalar completely transparent"""
from vedo import *

vol = Volume(dataurl+'embryo.slc').alpha([0,0,0.5]).c('k')

slices = []
for i in range(4):
    sl = vol.slice_plane(origin=[150,150,i*50+50], normal=(0,-1,1))
    slices.append(sl)

amap = [0, 1, 1, 1, 1]  # hide low value points giving them alpha 0
mslices = merge(slices) # merge all slices into a single Mesh
mslices.cmap('hot_r', alpha=amap).lighting('off').add_scalarbar3d()

show(vol, mslices, __doc__, axes=1).close()


"""RPC client tools"""
from __future__ import absolute_import

import os
import socket
import struct
import time

from . import base
from ..contrib import util
from .._ffi.base import TVMError
from .._ffi import function as function
from .._ffi import ndarray as nd
from ..module import load as _load_module

class RPCSession(object):
    """RPC Client session module

    Do not directly create the obhect, call connect
    """
    # pylint: disable=invalid-name
    def __init__(self, sess):
        self._sess = sess
        self._tbl_index = base._SessTableIndex(sess)
        self._remote_funcs = {}

    def get_function(self, name):
        """Get function from the session.

        Parameters
        ----------
        name : str
            The name of the function

        Returns
        -------
        f : Function
            The result function.
        """
        return self._sess.get_function(name)

    def context(self, dev_type, dev_id=0):
        """Construct a remote context.

        Parameters
        ----------
        dev_type: int or str

        dev_id: int, optional

        Returns
        -------
        ctx: TVMContext
            The corresponding encoded remote context.
        """
        ctx = nd.context(dev_type, dev_id)
        encode = (self._tbl_index + 1) * base.RPC_SESS_MASK
        ctx.device_type += encode
        ctx._rpc_sess = self
        return ctx

    def upload(self, data, target=None):
        """Upload file to remote runtime temp folder

        Parameters
        ----------
        data : str or bytearray
            The file name or binary in local to upload.

        target : str, optional
            The path in remote
        """
        if isinstance(data, bytearray):
            if not target:
                raise ValueError("target must present when file is a bytearray")
            blob = data
        else:
            blob = bytearray(open(data, "rb").read())
            if not target:
                target = os.path.basename(data)

        if "upload" not in self._remote_funcs:
            self._remote_funcs["upload"] = self.get_function(
                "tvm.rpc.server.upload")
        self._remote_funcs["upload"](target, blob)

    def download(self, path):
        """Download file from remote temp folder.

        Parameters
        ----------
        path : str
            The relative location to remote temp folder.

        Returns
        -------
        blob : bytearray
            The result blob from the file.
        """
        if "download" not in self._remote_funcs:
            self._remote_funcs["download"] = self.get_function(
                "tvm.rpc.server.download")
        return self._remote_funcs["download"](path)

    def load_module(self, path):
        """Load a remote module, the file need to be uploaded first.

        Parameters
        ----------
        path : str
            The relative location to remote temp folder.

        Returns
        -------
        m : Module
            The remote module containing remote function.
        """
        return base._LoadRemoteModule(self._sess, path)

    def cpu(self, dev_id=0):
        """Construct CPU device."""
        return self.context(1, dev_id)

    def gpu(self, dev_id=0):
        """Construct GPU device."""
        return self.context(2, dev_id)

    def cl(self, dev_id=0):
        """Construct OpenCL device."""
        return self.context(4, dev_id)

    def vulkan(self, dev_id=0):
        """Construct Vulkan device."""
        return self.context(7, dev_id)

    def metal(self, dev_id=0):
        """Construct Metal device."""
        return self.context(8, dev_id)

    def opengl(self, dev_id=0):
        """Construct OpenGL device."""
        return self.context(11, dev_id)

    def ext_dev(self, dev_id=0):
        """Construct extension device."""
        return self.context(12, dev_id)

class LocalSession(RPCSession):
    """RPCSession interface backed by local environment.

    This class can be used to implement functions that
    need to be ran both locally and remotely.
    """
    def __init__(self):
        # pylint: disable=super-init-not-called
        self.context = nd.context
        self.get_function = function.get_global_func
        self._temp = util.tempdir()

    def upload(self, data, target=None):
        if isinstance(data, bytearray):
            if not target:
                raise ValueError("target must present when file is a bytearray")
            blob = data
        else:
            blob = bytearray(open(data, "rb").read())
            if not target:
                target = os.path.basename(data)
        with open(self._temp.relpath(target), "wb") as f:
            f.write(blob)

    def download(self, path):
        return bytearray(open(self._temp.relpath(path), "rb").read())

    def load_module(self, path):
        return _load_module(self._temp.relpath(path))

class TrackerSession(object):
    """Tracker client session.

    Parameters
    ----------
    addr : tuple
        The address tuple
    """
    def __init__(self, addr):
        self._addr = addr
        self._sock = None
        self._connect()

    def __del__(self):
        self.close()

    def _connect(self):
        self._sock = base.connect_with_retry(self._addr)
        self._sock.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
        magic = struct.unpack("<i", base.recvall(self._sock, 4))[0]
        if magic != base.RPC_TRACKER_MAGIC:
            raise RuntimeError("%s is not RPC Tracker" % str(self._addr))

    def close(self):
        """Close the tracker connection."""
        if self._sock:
            self._sock.close()
            self._sock = None

    def summary(self):
        """Get the summary dict of the tracker."""
        base.sendjson(self._sock, [base.TrackerCode.SUMMARY])
        value = base.recvjson(self._sock)
        if value[0] != base.TrackerCode.SUCCESS:
            raise RuntimeError("Invalid return value %s" % str(value))
        return value[1]

    def text_summary(self):
        """Get a text summary of the tracker."""
        data = self.summary()

        total_ct = {}

        res = ""
        res += "Server List\n"
        res += "----------------------------\n"
        res += "server-address\tkey\n"
        res += "----------------------------\n"
        for item in data["server_info"]:
            addr = item["addr"]
            res += addr[0] + ":" + str(addr[1]) + "\t"
            res += item["key"] + "\n"
            key = item['key'].split(':')[1]   # 'server:rasp3b` -> 'rasp3b'
            if key not in total_ct:
                total_ct[key] = 0
            total_ct[key] += 1
        res += "----------------------------\n"
        res += "\n"

        # compute max length of device key
        queue_info = data['queue_info']
        keys = list(queue_info.keys())
        if keys:
            keys.sort()
            max_key_len = max([len(k) for k in keys])
        else:
            max_key_len = 0

        res += "Queue Status\n"
        title = ("%%-%ds" % max_key_len + "   total  free  pending\n") % 'key'
        separate_line = '-' * len(title) + '\n'
        res += separate_line + title + separate_line
        for k in keys:
            total = total_ct.get(k, 0)
            free, pending = queue_info[k]["free"], queue_info[k]["pending"]
            if total or pending:
                res += ("%%-%ds" % max_key_len + "   %-5d  %-4d  %-7d\n") % \
                       (k, total, free, pending)
        res += separate_line
        return res

    def request(self, key, priority=1, session_timeout=0, max_retry=5):
        """Request a new connection from the tracker.

        Parameters
        ----------
        key : str
            The type key of the device.

        priority : int, optional
            The priority of the request.

        session_timeout : float, optional
            The duration of the session, allows server to kill
            the connection when duration is longer than this value.
            When duration is zero, it means the request must always be kept alive.

        max_retry : int, optional
            Maximum number of times to retry before give up.
        """
        last_err = None
        for _ in range(max_retry):
            try:
                if self._sock is None:
                    self._connect()
                base.sendjson(self._sock,
                              [base.TrackerCode.REQUEST, key, "", priority])
                value = base.recvjson(self._sock)
                if value[0] != base.TrackerCode.SUCCESS:
                    raise RuntimeError("Invalid return value %s" % str(value))
                url, port, matchkey = value[1]
                return connect(url, port, matchkey, session_timeout)
            except socket.error as err:
                self.close()
                last_err = err
            except TVMError as err:
                last_err = err
        raise RuntimeError(
            "Cannot request %s after %d retry, last_error:%s" % (
                key, max_retry, str(last_err)))

    def request_and_run(self,
                        key,
                        func,
                        priority=1,
                        session_timeout=0,
                        max_retry=2):
        """Request a resource from tracker and run the func.

        This function safe-guard rare server node dropout during execution.
        In such case, a new resource will be requested and func will be ran again.

        Parameters
        ----------
        key : str
            The type key of the device.

        func : function of session -> value
            A stateless function

        priority : int, optional
            The priority of the request.

        session_timeout : float, optional
            The duration of the session, allows server to kill
            the connection when duration is longer than this value.
            When duration is zero, it means the request must always be kept alive.

        max_retry : int, optional
            Maximum number of times to retry the function before give up.
        """
        last_err = None
        for _ in range(max_retry):
            try:
                sess = self.request(key,
                                    priority=priority,
                                    session_timeout=session_timeout)
                tstart = time.time()
                return func(sess)
            except TVMError as err:
                duration = time.time() - tstart
                # roughly estimate if the error is due to timeout termination
                if session_timeout and duration >= session_timeout * 0.95:
                    raise RuntimeError(
                        "Session timeout when running %s" % func.__name__)
                last_err = err
        raise RuntimeError(
            "Failed to run on %s after %d retry, last_error:%s" % (
                key, max_retry, str(last_err)))

def connect(url, port, key="", session_timeout=0):
    """Connect to RPC Server

    Parameters
    ----------
    url : str
        The url of the host

    port : int
        The port to connect to

    key : str, optional
        Additional key to match server

    session_timeout : float, optional
        The duration of the session, allows server to kill
        the connection when duration is longer than this value.
        When duration is zero, it means the request must always be kept alive.

    Returns
    -------
    sess : RPCSession
        The connected session.
    """
    try:
        if session_timeout:
            key += " -timeout=%s" % str(session_timeout)
        sess = base._Connect(url, port, key)
    except NameError:
        raise RuntimeError("Please compile with USE_RPC=1")
    return RPCSession(sess)

def connect_tracker(url, port):
    """Connect to a RPC tracker

    Parameters
    ----------
    url : str
        The url of the host

    port : int
        The port to connect to

    Returns
    -------
    sess : TrackerSession
        The connected tracker session.
    """
    return TrackerSession((url, port))


"""
Runs a full AC continuation power flow.
"""
# --- run cpf ---
import logging

import numpy as np
from numpy import flatnonzero as find

import scipy.sparse as sp  # NOQA

from andes.shared import deg2rad  # NOQA
from andes.utils.misc import elapsed  # NOQA

import ams.pypower.utils as putil  # NOQA
import ams.pypower.io as pio  # NOQA
import ams.pypower.routines.opffcns as opfcn  # NOQA
from ams.pypower.idx import IDX  # NOQA
from ams.pypower.make import (makeSbus, makeYbus, dSbus_dV)  # NOQA
from ams.pypower.routines.pflow import newtonpf, pfsoln  # NOQA
from ams.pypower.core import ppoption  # NOQA
import ams.pypower.routines.cpf_callbacks as cpf_callbacks  # NOQA

logger = logging.getLogger(__name__)

def runcpf(casedata, ppopt=None, scale=1.2):
    """
    Runs a full AC continuation power flow.
    """
    ppopt = ppoption(ppopt)

    # options
    verbose = ppopt["VERBOSE"]
    step = ppopt["CPF_STEP"]
    parameterization = ppopt["CPF_PARAMETERIZATION"]
    adapt_step = ppopt["CPF_ADAPT_STEP"]
    cb_args = ppopt["CPF_USER_CALLBACK_ARGS"]

    # set up callbacks
    callback_names = ["cpf_default_callback"]
    if len(ppopt["CPF_USER_CALLBACK"]) > 0:
        if isinstance(ppopt["CPF_USER_CALLBACK"], list):
            callback_names = np.r_[callback_names, ppopt["CPF_USER_CALLBACK"]]
        else:
            callback_names.append(ppopt["CPF_USER_CALLBACK"])
    callbacks = []
    for callback_name in callback_names:
        callbacks.append(getattr(cpf_callbacks, callback_name))

    # read base case data
    ppcbase = pio.loadcase(casedata)
    nb = ppcbase["bus"].shape[0]

    # add zero columns to branch for flows if needed
    if ppcbase["branch"].shape[1] < IDX.branch.QT:
        ppcbase["branch"] = np.c_[ppcbase["branch"],
                                  np.zeros((ppcbase["branch"].shape[0],
                                            IDX.branch.QT - ppcbase["branch"].shape[1] + 1))]

    # convert to internal indexing
    ppcbase = opfcn.ext2int(ppcbase)
    baseMVAb, busb, genb, branchb = \
        ppcbase["baseMVA"], ppcbase["bus"], ppcbase["gen"], ppcbase["branch"]

    # get bus index lists of each type of bus
    ref, pv, pq = putil.bustypes(busb, genb)

    # generator info
    onb = find(genb[:, IDX.gen.GEN_STATUS] > 0)  # which generators are on?
    gbusb = genb[onb, IDX.gen.GEN_BUS].astype(int)  # what buses are they at?

    # scale the load and generation of base case as target case
    ppctarget = pio.loadcase(casedata)
    ppctarget["bus"][:, IDX.bus.PD] *= scale
    ppctarget["bus"][:, IDX.bus.QD] *= scale
    ppctarget["gen"][:, IDX.gen.PG] *= scale
    ppctarget["gen"][:, IDX.gen.QG] *= scale

    # add zero columns to branch for flows if needed
    if ppctarget["branch"].shape[1] < IDX.branch.QT:
        ppctarget["branch"] = np.c_[ppctarget["branch"],
                                    np.zeros((ppctarget["branch"].shape[0],
                                              IDX.branch.QT - ppctarget["branch"].shape[1] + 1))]

    # convert to internal indexing
    ppctarget = opfcn.ext2int(ppctarget)
    baseMVAt, bust, gent, brancht = \
        ppctarget["baseMVA"], ppctarget["bus"], ppctarget["gen"], ppctarget["branch"]

    # get bus index lists of each type of bus
    # ref, pv, pq = putil.bustypes(bust, gent)

    # generator info
    ont = find(gent[:, IDX.gen.GEN_STATUS] > 0)  # which generators are on?
    # gbust = gent[ont, IDX.gen.GEN_BUS].astype(int)  # what buses are they at?

    # -----  run the power flow  -----
    t0, _ = elapsed()  # start timer

    # initial state
    # V0    = np.ones((bus.shape[0])            ## flat start
    V0 = busb[:, IDX.bus.VM] * np.exp((1j * deg2rad * busb[:, IDX.bus.VA]))
    vcb = np.ones(V0.shape)    # create mask of voltage-controlled buses
    vcb[pq] = 0     # exclude PQ buses
    k = find(vcb[gbusb])    # in-service gens at v-c buses
    V0[gbusb[k]] = genb[onb[k], IDX.gen.VG] / abs(V0[gbusb[k]]) * V0[gbusb[k]]

    # build admittance matrices
    Ybus, Yf, Yt = makeYbus(baseMVAb, busb, branchb)

    # compute base case complex bus power injections (generation - load)
    Sbusb = makeSbus(baseMVAb, busb, genb)
    # compute target case complex bus power injections (generation - load)
    Sbust = makeSbus(baseMVAt, bust, gent)

    # scheduled transfer
    Sxfr = Sbust - Sbusb

    # Run the base case power flow solution
    lam = 0
    logger.info('Solve base case power flow')
    V, success, iterations = newtonpf(Ybus, Sbusb, V0, ref, pv, pq, ppopt)

    logger.info('Start full AC CPF')
    logger.info('%2d: lambda = %6.3f, %2d Newton steps' % (0, 0, iterations))

    lamprv = lam    # lam at previous step
    Vprv = V    # V at previous step
    continuation = 1
    cont_steps = 0

    # input args for callbacks
    cb_data = dict(ppc_base=ppcbase, ppc_target=ppctarget,
                   Sxfr=Sxfr, Ybus=Ybus, Yf=Yf, Yt=Yt,
                   ref=ref, pv=pv, pq=pq, ppopt=ppopt)
    cb_state = {}

    # invoke callbacks
    for k in range(len(callbacks)):
        cb_state, _ = callbacks[k](cont_steps, V, lam, V, lam,
                                   cb_data, cb_state, cb_args)

    if np.linalg.norm(Sxfr) == 0:
        logger.error('base case and target case have identical load and generation!')

        continuation = 0
        V0 = V
        lam0 = lam

    # tangent predictor z = [dx;dlam]
    z = np.zeros(2*len(V)+1)
    z[-1] = 1.0
    while continuation:
        cont_steps = cont_steps + 1
        # prediction for next step
        V0, lam0, z = cpf_predictor(V, lam, Ybus, Sxfr, pv, pq, step, z,
                                    Vprv, lamprv, parameterization)

        # save previous voltage, lambda before updating
        Vprv = V
        lamprv = lam

        # correction
        V, success, i, lam = cpf_corrector(Ybus, Sbusb, V0, ref, pv, pq,
                                           lam0, Sxfr, Vprv, lamprv, z,
                                           step, parameterization, ppopt)

        if not success:
            continuation = 0
            logger.warning('%2d: lambda = %6.3f, corrector did not converge in %d iterations' % (
                cont_steps, lam, i))
            break

        logger.info('%2d: lambda = %6.3f, %2d corrector Newton steps' %
                    (cont_steps, lam, i))

        # invoke callbacks
        for k in range(len(callbacks)):
            cb_state, _ = callbacks[k](cont_steps, V, lam, V0, lam0,
                                       cb_data, cb_state, cb_args)

        if isinstance(ppopt["CPF_STOP_AT"], str):
            if ppopt["CPF_STOP_AT"].upper() == "FULL":
                if abs(lam) < 1e-8:     # traced the full continuation curve
                    logger.info('Reached steady state loading limit in %d continuation steps' % cont_steps)
                    continuation = 0
                elif lam < lamprv and lam - step < 0:    # next step will overshoot
                    step = lam      # modify step-size
                    parameterization = 1    # change to natural parameterization
                    adapt_step = False      # disable step-adaptivity

            else:   # == 'NOSE'
                if lam < lamprv:    # reached the nose point
                    logger.info('Reached steady state loading limit in %d continuation steps' % cont_steps)
                    continuation = 0

        else:
            if lam < lamprv:
                logger.info('Reached steady state loading limit in %d continuation steps' % cont_steps)
                continuation = 0
            elif abs(ppopt["CPF_STOP_AT"] - lam) < 1e-8:     # reached desired lambda
                logger.info('Reached desired lambda %3.2f in %d continuation steps' % (
                    ppopt["CPF_STOP_AT"], cont_steps))
                continuation = 0
            # will reach desired lambda in next step
            elif lam + step > ppopt["CPF_STOP_AT"]:
                step = ppopt["CPF_STOP_AT"] - lam   # modify step-size
                parameterization = 1    # change to natural parameterization
                adapt_step = False      # disable step-adaptivity

        if adapt_step and continuation:
            pvpq = np.r_[pv, pq]
            # Adapt stepsize
            cpf_error = np.linalg.norm(np.r_[np.angle(V[pq]), abs(
                V[pvpq]), lam] - np.r_[np.angle(V0[pq]), abs(V0[pvpq]), lam0], np.Inf)
            if cpf_error < ppopt["CPF_ERROR_TOL"]:
                # Increase stepsize
                step = step * ppopt["CPF_ERROR_TOL"] / cpf_error
                if step > ppopt["CPF_STEP_MAX"]:
                    step = ppopt["CPF_STEP_MAX"]
            else:
                # decrese stepsize
                step = step * ppopt["CPF_ERROR_TOL"] / cpf_error
                if step < ppopt["CPF_STEP_MIN"]:
                    step = ppopt["CPF_STEP_MIN"]

    # invoke callbacks
    if success:
        cpf_results = {}
        for k in range(len(callbacks)):
            cb_state, cpf_results = callbacks[k](cont_steps, V, lam, V0, lam0,
                                                 cb_data, cb_state, cb_args, results=cpf_results, is_final=True)
    else:
        cpf_results = {}
        cpf_results["iterations"] = i

    # update bus and gen matrices to reflect the loading and generation
    # at the noise point
    bust[:, IDX.bus.PD] = busb[:, IDX.bus.PD] + lam * (bust[:, IDX.bus.PD] - busb[:, IDX.bus.PD])
    bust[:, IDX.bus.QD] = busb[:, IDX.bus.QD] + lam * (bust[:, IDX.bus.QD] - busb[:, IDX.bus.QD])
    gent[:, IDX.gen.PG] = genb[:, IDX.gen.PG] + lam * (gent[:, IDX.gen.PG] - genb[:, IDX.gen.PG])

    # update data matrices with solution
    bust, gent, brancht = pfsoln(baseMVAt, bust, gent, brancht,
                                 Ybus, Yf, Yt, V, ref, pv, pq)

    ppctarget["et"] = elapsed(t0)
    ppctarget["success"] = success

    # -----  output results  -----
    # convert back to original bus numbering & print results
    ppctarget["bus"], ppctarget["gen"], ppctarget["branch"] = bust, gent, brancht
    if success:
        n = cpf_results["iterations"] + 1
        cpf_results["V_p"] = opfcn.i2e_data(
            ppctarget, cpf_results["V_p"], np.full((nb, n), np.NaN), "bus", 0)
        cpf_results["V_c"] = opfcn.i2e_data(
            ppctarget, cpf_results["V_c"], np.full((nb, n), np.NaN), "bus", 0)
    results = opfcn.int2ext(ppctarget)
    results["cpf"] = cpf_results

    # zero out result fields of out-of-service gens & branches
    if len(results["order"]["gen"]["status"]["off"]) > 0:
        results["gen"][np.ix_(results["order"]["gen"]
                              ["status"]["off"], [IDX.gen.QF, IDX.gen.QG])] = 0

    if len(results["order"]["branch"]["status"]["off"]) > 0:
        results["branch"][np.ix_(results["order"]["branch"]
                                 ["status"]["off"], [IDX.branch.PF, IDX.branch.QF, IDX.branch.PT, IDX.branch.QT])] = 0

    if ppopt["CPF_PLOT_LEVEL"]:
        import matplotlib.pyplot as plt
        plt.show()
    sstats = dict(solver_name='PYPOWER', num_iters=cpf_results["iterations"])
    return results, success, sstats

def cpf_predictor(V, lam, Ybus, Sxfr, pv, pq,
                  step, z, Vprv, lamprv, parameterization):
    # sizes
    pvpq = np.r_[pv, pq]
    nb = len(V)
    npv = len(pv)
    npq = len(pq)

    # compute Jacobian for the power flow equations
    dSbus_dVm, dSbus_dVa = dSbus_dV(Ybus, V)

    j11 = dSbus_dVa[np.array([pvpq]).T, pvpq].real
    j12 = dSbus_dVm[np.array([pvpq]).T, pq].real
    j21 = dSbus_dVa[np.array([pq]).T, pvpq].imag
    j22 = dSbus_dVm[np.array([pq]).T, pq].imag

    J = sp.vstack([
        sp.hstack([j11, j12]),
        sp.hstack([j21, j22])
    ], format="csr")

    dF_dlam = -np.r_[Sxfr[pvpq].real, Sxfr[pq].imag].reshape((-1, 1))
    dP_dV, dP_dlam = cpf_p_jac(parameterization, z, V, lam, Vprv, lamprv, pv, pq)

    # linear operator for computing the tangent predictor
    J = sp.vstack([
        sp.hstack([J, dF_dlam]),
        sp.hstack([dP_dV, dP_dlam])
    ], format="csr")

    Vaprv = np.angle(V)
    Vmprv = abs(V)

    # compute normalized tangent predictor
    s = np.zeros(npv+2*npq+1)
    s[-1] = 1
    z[np.r_[pvpq, nb+pq, 2*nb]] = sp.linalg.spsolve(J, s)
    z = z / np.linalg.norm(z)

    Va0 = Vaprv
    Vm0 = Vmprv
    lam0 = lam

    # prediction for next step
    Va0[pvpq] = Vaprv[pvpq] + step * z[pvpq]
    Vm0[pq] = Vmprv[pq] + step * z[nb+pq]
    lam0 = lam + step * z[2*nb]
    V0 = Vm0 * np.exp((1j * Va0))

    return V0, lam0, z

def cpf_corrector(Ybus, Sbus, V0, ref, pv, pq,
                  lam0, Sxfr, Vprv, lamprv, z,
                  step, parameterization, ppopt):

    # default arguments
    if ppopt is None:
        ppopt = ppoption(ppopt)

    # options
    verbose = ppopt["VERBOSE"]
    tol = ppopt["PF_TOL"]
    max_it = ppopt["PF_MAX_IT"]

    # initialize
    converged = 0
    i = 0
    V = V0
    Va = np.angle(V)
    Vm = abs(V)
    lam = lam0

    # set up indexing for updating V
    pvpq = np.r_[pv, pq]
    npv = len(pv)
    npq = len(pq)
    nb = len(V)
    j1 = 0
    j2 = npv    # j1:j2 - V angle of pv buses
    j3 = j2
    j4 = j2 + npq   # j1:j2 - V angle of pv buses
    j5 = j4
    j6 = j4 + npq   # j5:j6 - V mag of pq buses
    j7 = j6
    j8 = j6 + 1    # j7:j8 - lambda

    # evaluate F(x0, lam0), including Sxfr transfer/loading

    mis = V * np.conj(Ybus.dot(V)) - Sbus - lam*Sxfr
    F = np.r_[mis[pvpq].real, mis[pq].imag]

    # evaluate P(x0, lambda0)
    P = cpf_p(parameterization, step, z, V, lam, Vprv, lamprv, pv, pq)

    # augment F(x, lambda) with P(x, lambda)
    F = np.r_[F, P]

    # check tolerance
    normF = np.linalg.norm(F, np.Inf)
    logger.debug('CPF correction')
    logger.debug('%2d: |F(x)| = %.10g', i, normF)
    if normF < tol:
        converged = 1

    # do Newton iterations
    while not converged and i < max_it:
        # update iteration counter
        i = i + 1

        # evaluate Jacobian
        dSbus_dVm, dSbus_dVa = dSbus_dV(Ybus, V)

        j11 = dSbus_dVa[np.array([pvpq]).T, pvpq].real
        j12 = dSbus_dVm[np.array([pvpq]).T, pq].real
        j21 = dSbus_dVa[np.array([pq]).T, pvpq].imag
        j22 = dSbus_dVm[np.array([pq]).T, pq].imag

        J = sp.vstack([
            sp.hstack([j11, j12]),
            sp.hstack([j21, j22])
        ], format="csr")

        dF_dlam = -np.r_[Sxfr[pvpq].real, Sxfr[pq].imag].reshape((-1, 1))
        dP_dV, dP_dlam = cpf_p_jac(parameterization, z, V, lam, Vprv, lamprv, pv, pq)

        # augment J with real/imag -Sxfr and z^T
        J = sp.vstack([
            sp.hstack([J, dF_dlam]),
            sp.hstack([dP_dV, dP_dlam])
        ], format="csr")

        # compute update step
        dx = -1 * sp.linalg.spsolve(J, F)

        # update voltage
        if npv:
            Va[pv] = Va[pv] + dx[j1:j2]
        if npq:
            Va[pq] = Va[pq] + dx[j3:j4]
            Vm[pq] = Vm[pq] + dx[j5:j6]
        V = Vm * np.exp((1j * Va))
        Vm = abs(V)
        Va = np.angle(V)

        # update lambda
        lam = lam + dx[j7:j8]

        # evalute F(x, lam)
        mis = V * np.conj(Ybus.dot(V)) - Sbus - lam*Sxfr
        F = np.r_[mis[pv].real, mis[pq].real, mis[pq].imag]

        # evaluate P(x, lambda)
        P = cpf_p(parameterization, step, z, V, lam, Vprv, lamprv, pv, pq)

        # augment F(x, lambda) with P(x, lambda)
        F = np.r_[F, P]

        # check for convergence
        normF = np.linalg.norm(F, np.Inf)
        logger.debug('%2d: |F(x)| = %.10g', i, normF)

        if normF < tol:
            converged = 1

    if not converged:
        logger.error('Newton''s method corrector did not converge in %d iterations.' % i)

    return V, converged, i, lam

def cpf_p_jac(parameterization, z, V, lam, Vprv, lamprv, pv, pq):
    if parameterization == 1:
        npv = len(pv)
        npq = len(pq)
        dP_dV = np.zeros(npv+2*npq)
        if lam >= lamprv:
            dP_dlam = 1.0
        else:
            dP_dlam = -1.0

    elif parameterization == 2:
        pvpq = np.r_[pv, pq]

        Va = np.angle(V)
        Vm = abs(V)
        Vaprv = np.angle(Vprv)
        Vmprv = abs(Vprv)
        dP_dV = 2 * (np.r_[Va[pvpq], Vm[pq]] -
                     np.r_[Vaprv[pvpq], Vmprv[pq]])
        if lam == lamprv:
            dP_dlam = 1.0
        else:
            dP_dlam = 2 * (lam - lamprv)

    elif parameterization == 3:
        nb = len(V)
        dP_dV = z[np.r_[pv, pq, nb+pq]]
        dP_dlam = z[2 * nb]

    return dP_dV, dP_dlam

def cpf_p(parameterization, step, z, V, lam, Vprv, lamprv, pv, pq):
    # evaluate P(x0, lambda0)
    if parameterization == 1:
        if lam >= lamprv:
            P = lam - lamprv - step
        else:
            P = lamprv - lam - step

    elif parameterization == 2:
        pvpq = np.r_[pv, pq]

        Va = np.angle(V)
        Vm = abs(V)
        Vaprv = np.angle(Vprv)
        Vmprv = abs(Vprv)
        P = sum(np.square(np.r_[Va[pvpq], Vm[pq], lam] -
                          np.r_[Vaprv[pvpq], Vmprv[pq], lamprv])) - np.square(step)

    elif parameterization == 3:
        pvpq = np.r_[pv, pq]

        nb = len(V)
        Va = np.angle(V)
        Vm = abs(V)
        Vaprv = np.angle(Vprv)
        Vmprv = abs(Vprv)
        P = np.dot(z[np.r_[pvpq, nb+pq, 2*nb]].reshape((1, -1)),
                   np.array(np.r_[Va[pvpq], Vm[pq], lam] - np.r_[Vaprv[pvpq], Vmprv[pq], lamprv]).T) - step

    return P


import argparse
import concurrent.futures
import itertools
import logging
import pathlib
import sys

import numpy as np

from . import data
from . import utils

logger = logging.getLogger(__name__)

def _do_convert(args):
    mesh = data.load_mesh(args.input)
    data.save_mesh(args.output, mesh)

def _do_shoot(args):
    mesh = data.load_mesh(str(args.input))

    has_min_holes = args.min_holes is not None
    has_max_holes = args.max_holes is not None
    if has_min_holes != has_max_holes:
        raise ValueError("--min-holes and --max-holes must be set together")
    n_holes = ((args.min_holes, args.max_holes)
               if has_min_holes
               else args.holes)

    has_min_dropout = args.min_dropout is not None
    has_max_dropout = args.max_dropout is not None
    if has_min_dropout != has_max_dropout:
        raise ValueError(
            "--min-dropout and --max-dropout must be set together")
    dropout = ((args.min_dropout, args.max_dropout)
               if has_min_dropout
               else args.dropout)

    mask_faces = (np.load(args.mask) if args.mask is not None
                  else None)
    faces = None if mask_faces is None else mesh.faces
    point_indices = utils.shoot_holes(mesh.vertices,
                                      n_holes,
                                      dropout,
                                      mask_faces=mask_faces,
                                      faces=faces)
    shot = utils.remove_points(mesh, point_indices)
    shot.save(str(args.output))

def identify_meshes(dir_):
    """List meshes and identify the challenge/track in a directory tree."""
    meshes_track1 = list(dir_.glob("**/*_normalized.npz"))
    meshes_track2 = list(dir_.glob("**/fusion_textured.npz"))
    meshes_challenge2 = list(dir_.glob("**/model_*.obj"))
    if meshes_track1:
        meshes = sorted(meshes_track1)
        challenge = 1
        track = 1
    elif meshes_track2:
        meshes = sorted(meshes_track2)
        challenge = 1
        track = 2
    elif meshes_challenge2:
        meshes = sorted(meshes_challenge2)
        challenge = 2
        track = None
    else:
        meshes = []
        challenge = None
        track = None

    return meshes, challenge, track

def shoot_helper(mesh_index,
                 path,
                 n_holes,
                 dropout,
                 shape_seeds,
                 n_meshes,
                 input_dir,
                 output_dir,
                 mask_dir=None,
                 ):
    logger.info(f"{mesh_index + 1}/{n_meshes} -------- TEST ------------ processing {path}")

    rel_path = path.relative_to(input_dir)

    # Lazily load the mesh and mask when it is sure they are needed.
    mesh_cached = None
    mask_cached = None

    def make_name_suffix(shape_index, n_shapes):
        if n_shapes == 1:
            return "partial"
        else:
            # Assuming 0 <= shape_index <= 99.
            return f"partial-{shape_index:02d}"

    def load_mask(rel_path):
        mask_name = f"{rel_path.stem}-mask.npy"
        mask_rel_path = rel_path.with_name(mask_name)
        mask_path = mask_dir / mask_rel_path
        logger.info(f"loading mask {mask_path}")
        mask_faces = np.load(mask_path)
        return mask_faces

    n_shapes = len(shape_seeds)
    for shape_index, shape_seed in enumerate(shape_seeds):
        out_name_suffix = make_name_suffix(shape_index, n_shapes)
        out_name = f"{rel_path.stem}-{out_name_suffix}.npz"
        out_rel_path = rel_path.with_name(out_name)
        out_path = output_dir / out_rel_path

        if out_path.exists():
            logger.warning(f"shape exists, skipping {out_path}")
            continue

        logger.info(f"generating shape {shape_index + 1}/{n_shapes}")

        out_path.parent.mkdir(parents=True, exist_ok=True)

        if mesh_cached is None:
            mesh_cached = data.load_mesh(str(path))
        if mask_cached is None and mask_dir is not None:
            mask_cached = load_mask(rel_path)

        mesh = mesh_cached
        mask = mask_cached

        logger.info(f"TEST shape seed = {shape_seed}")
        shape_rng = np.random.default_rng(shape_seed)
        point_indices = utils.shoot_holes(mesh.vertices,
                                          n_holes,
                                          dropout,
                                          mask_faces=mask,
                                          faces=mesh.faces,
                                          rng=shape_rng)
        partial = utils.remove_points(mesh, point_indices)

        logger.info(f"{shape_index + 1}/{n_shapes} saving {out_path}")
        partial.save(str(out_path))

def _do_shoot_dir(args):
    """Generate partial data in a directory tree of meshes.

    An independent PRNG is used for each partial shape. Independent seeds are
    created in advance for each partial shape. The process is thus reproducible
    and can also be interrupted and resumed without generating all the previous
    shapes.
    """
    input_dir = args.input_dir
    output_dir = args.output_dir
    mask_dir = args.mask_dir
    seed = args.seed
    n_holes = args.holes
    dropout = args.dropout
    n_shapes = args.n_shapes
    n_workers = args.n_workers

    logger.info("generating partial data in directory tree")
    logger.info(f"input dir = {input_dir}")
    logger.info(f"output dir = {output_dir}")
    logger.info(f"mask dir = {mask_dir}")
    logger.info(f"seed = {seed}")
    logger.info(f"holes = {n_holes}")
    logger.info(f"dropout = {dropout}")
    logger.info(f"n_shapes = {n_shapes}")
    logger.info(f"n_workers = {n_workers}")

    mesh_paths, challenge, track = identify_meshes(input_dir)
    if challenge is None:
        raise ValueError(f"could not identify meshes in {input_dir}")
    logger.info(f"detected challenge {challenge} track {track}")
    n_meshes = len(mesh_paths)
    logger.info(f"found {n_meshes} meshes")

    logger.info(f"setting random seed {seed}")
    rng = np.random.default_rng(seed)

    # Create in advance the individual initial random states for each partial
    # shape to be generated so that the process of generation can be
    # - resumed without regenerating the previous shapes,
    # - parallelised.
    seeds = rng.integers(1e12, size=(n_meshes, n_shapes))

    mesh_indices = range(n_meshes)
    with concurrent.futures.ProcessPoolExecutor(max_workers=n_workers) as executor:
        executor.map(
            shoot_helper,
            mesh_indices,
            mesh_paths,
            itertools.repeat(n_holes),
            itertools.repeat(dropout),
            seeds.tolist(),
            itertools.repeat(n_meshes),
            itertools.repeat(input_dir),
            itertools.repeat(output_dir),
            itertools.repeat(mask_dir),
        )

def _parse_args():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers()

    parser_convert = subparsers.add_parser(
        "convert",
        help="Convert between mesh formats.",
    )
    parser_convert.add_argument("input", type=pathlib.Path)
    parser_convert.add_argument("output", type=pathlib.Path)
    parser_convert.set_defaults(func=_do_convert)

    parser_shoot = subparsers.add_parser(
        "shoot",
        help="Generate partial data with the shooting method.",
    )
    parser_shoot.add_argument("input", type=pathlib.Path)
    parser_shoot.add_argument("output", type=pathlib.Path)
    parser_shoot.add_argument(
        "--holes", type=int, default=40,
        help="Number of holes to shoot.",
    )
    parser_shoot.add_argument(
        "--min-holes", type=int, default=None,
        help="Minimum number of holes to generate."
             " (Supersedes --holes and requires --max-holes.)",
    )
    parser_shoot.add_argument(
        "--max-holes", type=int, default=None,
        help="Maximum number of holes to generate."
             " (Supersedes --holes and requires --min-holes.)",
    )
    parser_shoot.add_argument(
        "--dropout", type=float, default=2e-2,
        help="Proportion of points of the mesh to remove in a single hole.",
    )
    parser_shoot.add_argument(
        "--min-dropout", type=float, default=None,
        help="Minimum proportion of points of the mesh to remove in a single "
             "hole."
             " (Supersedes --dropout and requires --max-dropout.)",
    )
    parser_shoot.add_argument(
        "--max-dropout", type=float, default=None,
        help="Maximum proportion of points of the mesh to remove in a single "
             "hole."
             " (Supersedes --dropout and requires --min-dropout.)",
    )
    parser_shoot.add_argument(
        "--mask", type=pathlib.Path,
        help=" (optional) Path to the mask (.npy) to generate holes only on"
             " regions considered for evaluation.",
    )
    parser_shoot.set_defaults(func=_do_shoot)

    parser_shoot_dir = subparsers.add_parser(
        "shoot_dir",
        help="Generate partial data with the shooting method for a directory"
             " tree of meshes."
    )
    parser_shoot_dir.add_argument("input_dir", type=pathlib.Path)
    parser_shoot_dir.add_argument("output_dir", type=pathlib.Path)
    parser_shoot_dir.add_argument(
        "--mask-dir",
        type=pathlib.Path,
        help=" (optional) Directory tree with the masks (.npy). If defined,"
             " the partial data is created only on the non-masked faces of the"
             " meshes. (Only valid for challenge 1.)",
    )
    parser_shoot_dir.add_argument(
        "--holes",
        type=int,
        default=40,
        help="Number of holes to shoot.",
    )
    parser_shoot_dir.add_argument(
        "--dropout",
        type=float,
        default=2e-2,
        help="Proportion of points of the mesh to remove in a single hole.",
    )
    parser_shoot_dir.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Initial state for the pseudo random number generator."
             " If not set, the initial state is not set explicitly.",
    )
    parser_shoot_dir.add_argument(
        "-n", "--n-shapes",
        type=int,
        default=1,
        help="Number of partial shapes to generate per mesh."
             " If n = 1 (default), the shape is saved as"
             " '<mesh_name>-partial.npz'."
             " If n > 1, the shapes are saved as '<mesh_name>-partial-XY.npz'"
             ", with XY as 00, 01, 02... (assuming n <= 99).",
    )
    parser_shoot_dir.add_argument(
        "--n-workers",
        type=int,
        default=None,
        help="Number of parallel processes. By default, the number of"
             " available processors.",
    )
    parser_shoot_dir.set_defaults(func=_do_shoot_dir)

    args = parser.parse_args()

    # Ensure the help message is displayed when no command is provided.
    if "func" not in args:
        parser.print_help()
        sys.exit(1)

    return args

def main():
    args = _parse_args()
    args.func(args)

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    main()


#
# Python implementation of the Connect Four-Game
#
# File: connect-four.py
# Date: 2023-06-21
#
# General idea: Connect Four can be understood as sequence of
# moves. These game sequence can be encoded as string where each
# character names the chosen columns (from 1 to 7).
#
# For implementation purposes this game sequence can be represented by
# a dictionary. Keys of the dictionary are the (resulting) position on
# the board as tuples with row and column (on a grid with 6 rows and 7
# columns).  Row and column in position tuples are zero-based. The
# values of such a dictionary hold the respective player index (0 or
# 1).
#
# Note: The `encode` function relies on the fact that Python
# dictionaries preserve insertion order (since version 3.7)
#

# CONSTANTS
MAXROWS = 6
MAXCOLS = 7
# ASCII encoded characters for the players
PLAYER1 = '\033[1;31;40mR\033[0;37;40m' # red on black
PLAYER2 = '\033[1;34;40mB\033[0;37;40m' # blue on black
# '\033[1;32;40mG\033[0;37;40m' # green on black
# '\033[1;33;40mG\033[0;37;40m' # yellow on black
PLAYERS = { 0 : PLAYER1, 1 : PLAYER2 }
# Default verbosity
VERBOSE = False

# ENCODING/DECODING
def decode(s, verbose=VERBOSE):
    """Return dictionary from game sequence `s` with positions (int, int)
    as key and player index (int). Decoding stops if resulting
    position is not valid, i.e. exceeds maximum number of rows or
    columns, and if a subsequence leads to a win.

    s: str
    Returns: dict of type { (int, int) : int }

    """
    dic = {} # resulting dict
    # local dict with free rows per columns (initially 0)
    row = { i : 0 for i in range(MAXCOLS) }
    if s == '':
        return dic
    else:
        for i in range(len(s)):
            c = int(s[i]) - 1 # column index zero based
            p = i % 2 # player index starting with 0
            r = row.get(c, -1) # free row in column c
            if is_win(dic):
                if verbose:
                    print("Sequence '%s' leads to win yet!" % s[:i])
                return dic
            elif 0 <= c < MAXCOLS and 0 <= r < MAXROWS:
                dic[(r, c)] = p # add pos and player into dict
                row[c] = r + 1 # increase row for column c
            else:
                if verbose:
                    print("Sequence '%s' leads to filled column %d yet!" % (s[:i], c+1))
                return dic
    return dic

def encode(d):
    """Return string with game sequence from dictionary `d`.

    d: dict
    Returns: str
    """
    return ''.join([str(k[1]+1) for k in d.keys()])

# PRINTABLE REPRESENTATION
def grid(s):
    """Returns printable string representing the board after game sequence
    `s` (string or dictionary). Moves are colored using ASCII
    encoding.

    s: str or dict of type { (int, int) : int }
    Returns: str
    """
    p = '' # the result string
    if isinstance(s, str):
        d = decode(s)
    elif isinstance(s, dict):
        d = s
    else:
        print("Argument '%s' is not a valid game sequence!" % s)
        return None

    for r in reversed(range(MAXROWS)):
        p += ' |'
        for c in range(MAXCOLS):
            if (r, c) in d:
                p += " " + PLAYERS[d[(r, c)]] + " "
            else:
                p += '   '
        p += '| ' + str(r+1) + '\n'
    b = '' # bottom line
    for i in range(MAXCOLS):
        b += '-' + str(i+1) + '-'
    h = '\n Next turn: ' + PLAYERS[len(d)%2]
    return p + ' +' + b + '+' + h

# TESTING
def is_win(d, verbose=VERBOSE):
    """Returns True if game dictionary d shows a win. Prints information
    about win if optional parameter `verbose` is set to `True`.

    d: dict
    verbose: bool (defaults to False)
    Returns: bool
    """
    if len(d) < 7: return False # number of moves to low

    # row, column and player indices of last move
    ((r, c), x) = list(d.items())[-1]

    # vertical
    b1 = (r-1,c) in d and d[(r-1,c)] == x
    b2 = (r-2,c) in d and d[(r-2,c)] == x
    b3 = (r-3,c) in d and d[(r-3,c)] == x
    if b1 and b2 and b3:
        if verbose:
            print("%s wins! Vertical at (%d, %d) [%s]." % (PLAYERS[x], r+1, c+1, encode(d)))
        return True
    # diagonal (upwards)
    b0 = (r+3,c+3) in d and d[(r+3,c+3)] == x
    b1 = (r+2,c+2) in d and d[(r+2,c+2)] == x
    b2 = (r+1,c+1) in d and d[(r+1,c+1)] == x
    b3 = (r-1,c-1) in d and d[(r-1,c-1)] == x
    b4 = (r-2,c-2) in d and d[(r-2,c-2)] == x
    b5 = (r-3,c-3) in d and d[(r-3,c-3)] == x
    if (b0 and b1 and b2) or (b1 and b2 and b3) or (b2 and b3 and b4) or (b3 and b4 and b5):
        if verbose:
            print("%s wins! Diagonal at (%d, %d) [%s]." % (PLAYERS[x], r+1, c+1, encode(d)))
        return True

    # diagonal (downwards)
    b0 = (r+3,c-3) in d and d[(r+3,c-3)] == x
    b1 = (r+2,c-2) in d and d[(r+2,c-2)] == x
    b2 = (r+1,c-1) in d and d[(r+1,c-1)] == x
    b3 = (r-1,c+1) in d and d[(r-1,c+1)] == x
    b4 = (r-2,c+2) in d and d[(r-2,c+2)] == x
    b5 = (r-3,c+3) in d and d[(r-3,c+3)] == x
    if (b0 and b1 and b2) or (b1 and b2 and b3) or (b2 and b3 and b4) or (b3 and b4 and b5):
        if verbose:
            print("%s wins! Diagonal at (%d, %d) [%s]." % (PLAYERS[x], r+1, c+1, encode(d)))
        return True

    # check horizontal (from left to right)
    b0 = (r,c-3) in d and d[(r,c-3)] == x
    b1 = (r,c-2) in d and d[(r,c-2)] == x
    b2 = (r,c-1) in d and d[(r,c-1)] == x
    b3 = (r,c+1) in d and d[(r,c+1)] == x
    b4 = (r,c+2) in d and d[(r,c+2)] == x
    b5 = (r,c+3) in d and d[(r,c+3)] == x
    if (b0 and b1 and b2) or (b1 and b2 and b3) or (b2 and b3 and b4) or (b3 and b4 and b5):
        if verbose:
            print("%s wins! Horizontal at (%d, %d) [%s]." % (PLAYERS[x], r+1, c+1, encode(d)))
        return True

    return False

def is_final(d, verbose=VERBOSE):
    """Returns `True` if game dictionary `d` is a win or no further move
    is possible because the board is filled.

    s: dict of type { (int, int) : int }
    Returns: bool
    """
    return is_win(d, verbose=verbose) or len(d) == MAXROWS * MAXCOLS

# GENERATORS
def next_moves(d):
    """Returns list with possible next tupels for game dictionary
    d. Tuples contain position and player. List may be empty if
    there is no next move.

    d: dict of type ((int, int), int)
    Returns: list of type ((int, int), int)
    """
    l = [] # result list
    p = len(d) % 2 # player for next move
    # iterate over columns
    for c in range(MAXCOLS):
        r = 0 # initial row
        # positions in column c?
        ps = filter(lambda t: t[1] == c, d.keys())
        for t in ps:
            r = max(t[0]+1, r)
        if r < MAXROWS:
            l.append(((r, c), p))
    return(l)

def next_dicts(d = {}, lim = 1, out=None, verbose=VERBOSE):
    """Returns list of all possible game dictionaries starting with d. The
    number of additional moves is limited by the optional
    parameter. Prints additional information about wins if optional
    paramter 'verbose' is set to 'True'. With optional parameter
    out=filename results are written to filname (in string encoding).

    d: dict of type ((int, int), int) (defaults to {})
    lim: int (defaults to 1)
    out: str (defaults to None)
    verbose: bool (defaults to True)

    Returns: list of dicts with type ((int, int), int)
    """
    ds = [d] # start list
    rs = [] # result list
    if not out == None:
        fn = open(out, 'w')

    ll = min(len(d) + lim, MAXROWS * MAXCOLS)
    while len(ds) > 0:
        x = ds.pop()
        f = is_final(x, verbose=verbose)
        w = is_win(x, verbose=verbose)
        if f or len(x) >= ll:
            rs.append(x)
            if not out == None:
                # Writes encoded grid, and array with is final, is win and last player number (1 or 2)
                fn.write("%s [final: %s, win: %s, player: %d]\n" % (encode(x), f, w, 2 - (len(x) % 2)))
        else:
            for (po, pl) in next_moves(x):
                dc = x.copy()
                dc[po] = pl
                ds.append(dc)
    if not out == None:
        fn.close
    return(rs)

# VALUATION
def best_moves(d = {}, lim=4):
    """Return a list of optimal next move(s) for next player in a game
    with dictionary `d`. The algorithm looks upto `lim` moves ahead
    and regards all possible wins and losses. The list elements are
    tuples of an integer score and list of moves with that score. The
    list is ordered by descending score. Higher scores indicate moves
    which either lead to a win or prevent losses.

    d: dict of type ((int, int), int)
    lim: int (defaults to 4)
    Returns: list of type (int, ((int, int), int))
    """
    length = len(d)
    # dict with moves and priorities
    moves = { m : 0 for m in next_moves(d) }
    # all next win dicts up to 'lim'
    dicts = [ n for n in next_dicts(d, lim=lim, verbose=False) if is_win(n) ]
    final = False # break condition
    ahead = 1 # look ahead
    while not final and ahead <= lim:
        # win dicts with len = look ahead
        dl = [ m for m in dicts if len(m)-length == ahead ]
        # print("Look ahead = %d, no. of games: %d" % (ahead, len(dl)))
        while len(dl) > 0:
            cd = dl.pop()
            cm = list(cd.items())[length] # current move
            if ahead % 2 == 1: # current player wins
                # result.append(encode(cd)[length:] + "+")
                moves[cm] += 1
            else: # current player loses
                # result.append(encode(cd)[length:] + "-")
                moves[cm] -= 1
            final = True
        ahead += 1
    # transpose moves dict
    rdict = { s : [] for s in sorted(moves.values(), reverse=True) }
    for (m, s) in moves.items():
        rdict[s].append(m)
    return list(rdict.items())

# PLAYING
def play(d = {}, auto=[]):
    """Starts an interactive game play starting with game dictionary
    `d`. The game can be terminated by entering 'q' or 'Q'. Moves for
    players in list `auto` are automatically chosen.

    d: dict of type {(int, int), int} (defaults to {})
    auto: list of type int (defaults to [])
    Returns: dict of type { (int, int): int }
    """
    while not is_final(d, verbose=True):
        print('\n' + str(grid(d)))
        p = len(d) % 2 # player index
        val = False # valid input
        brk = False
        nxt = next_moves(d)
        while not val:
            if p in auto:
                bst = best_moves(d)[0][1] # best_move(s)
                mov = bst[len(bst)//2][0] # median move
                ans = str(mov[1]+1)
            else:
                ans = input(' Enter column, Player %s: ' % PLAYERS[p])
            if ans in ['q', 'Q']:
                brk = True
                val = True
            elif ans in [str(i+1) for i in range(MAXCOLS)]:
                c = int(ans)-1
                m = next((m for m in nxt if m[0][1] == c), None)
                if m == None:
                    print(" Column '%s' not allowed!" % ans)
                else:
                    print(" Plays column '%s'!" % ans)
                    val = True
            else:
                print(" Invalid input '%s'!" % ans)
        if brk:
            break
        else:
            d[m[0]] = m[1]
    return d


#import pandas as pd
import re
import csv

""" sql files used for testing

sqlfile1 = "Sample-SQL-File-500000-Rows.sql"
test1 = "database/test.sql"
test2 = "database/test3.sql"
test3 = "database/test4.sql"
test_list = [sqlfile1, test1, test2, test3]
sqlfile = test_list[2]
"""

def sqlreader(sqlfile):
    with open(sqlfile) as sqlfile:
        reader = csv.reader(sqlfile, delimiter=' ', quotechar='|')
        sqldatafile = ''
        for row in reader:
            sqldatafile = sqldatafile + ''.join(row) + '\n'
    return sqldatafile

def tablesfinder(sqldatafile):
    #tables
    table_names = re.findall('CREATETABLE[\w]*`([A-Za-z0-9.:;!?<>()-+=\s\w\d]*)`\(', sqldatafile)
    if len(table_names) > 0:
        print('Tables names: "{}"'.format(table_names))
        print('Amount of Tables: {}'.format(len(table_names)))

        pattern = '|INSERTINTO'
        inserts = sqldatafile.replace('INSERTINTO',pattern).split('|')
        print('Inserts Amount: {}'.format(len(inserts)))
        i=-1
        inserts_table_names = []
        for t in inserts[1:]:
            i +=1
            #columns names and counting for table
            columns_names = re.search('INSERTINTO\W([A-Za-z0-9.:;!?<>\(\)-+=\s\w\d]*)\W*([\WA-Za-z0-9.:;!?<>()-+=\s\w\d]*)\WVALUES', t)
            if columns_names:
                print(columns_names[1])
                columns = columns_names[2].replace('`', '').split(',')
                columns_str = ''
                for column in columns:
                    columns_str = columns_str + '"' + str(column) + '", '
                print('Columns Names in Table "{}" is: {}'.format(columns_names[1], columns_str[:-2]))
                print('Amount of Columns in Table "{}" is: {}'.format(columns_names[1], len(columns)))
                inserts_table_names.append(columns_names[1])
            else:
                pass

        """ This piece of code was a try for rows count
        it stuck during multiple INSERTS between different tables names,
        but there is a ways to fix it if having a time and motivation for

        ii=-1
        insert_index = []
        for insert_name in inserts_table_names:
            ii+=1
            insert_index.append(ii)
        print(inserts_table_names)
        print(insert_index)

        iii=-1
        rows_sum = 0
        for insert_name in inserts_table_names:
            rows_sum = 0
            iii+=1
            if len(inserts_table_names) < 19:
                rows_sum = inserts[iii].count('\n')
                if insert_name != inserts_table_names[iii+1]:
                    rows_sum += inserts[iii+2].count('\n')
            else:
                rows_sum = inserts[ii+1].count('\n')
                print('Amount of Rows in Table "{}": {}\n'.format(insert_name, rows_sum))

        #rows counter for table

            #if columns_names[1] ==
            #rows_sum += insert.count('\n')
            #print('Amount of Rows in Table "{}": {}\n'.format(table_names[i], rows_sum))
            """

    else:
        print('No Tables detected\n')

def main():
    tablesfinder(sqlreader(sqlfile))

main()


# progarm that calculates the Fibonacci value after user input

print("This program that computes the nth Fibonacci number where the n is a value input by the user")
num = eval(input("how many numbers would you like to evalute? "))

def f(n,a = 0, b = 1):
    k = [a,b]
    while len(k) < n:
        k = k +([k[len(k)-1] + k[len(k)-2]])
    return k

b = sum(f(num + 1))
print(b)