import numpy as np
import logging
logger = logging.getLogger(name=__name__)

from ...sgmcmc_sampler import SGMCMCHelper
from .kernels import (
        LGSSMPriorKernel,
        LGSSMOptimalKernel,
        LGSSMHighDimOptimalKernel,
        )
from ...particle_filters.buffered_smoother import (
        buffered_pf_wrapper,
        average_statistic,
        )
from ..._utils import lower_tri_mat_inv

class LGSSMHelper(SGMCMCHelper):
    """ LGSSM Helper

        forward_message (dict) with keys
            log_constant (double) log scaling const
            mean_precision (ndarray) mean precision
            precision (ndarray) precision

        backward_message (dict) with keys
            log_constant (double) log scaling const
            mean_precision (ndarray) mean precision
            precision (ndarray) precision
    """
    def __init__(self, n, m, forward_message=None, backward_message=None,
            **kwargs):
        self.n = n
        self.m = m

        if forward_message is None:
             forward_message = {
                    'log_constant': 0.0,
                    'mean_precision': np.zeros(self.n),
                    'precision': np.eye(self.n)/10,
                    }
        self.default_forward_message=forward_message

        if backward_message is None:
            backward_message = {
                'log_constant': 0.0,
                'mean_precision': np.zeros(self.n),
                'precision': np.zeros((self.n, self.n)),
                }
        self.default_backward_message=backward_message
        return

    ## Message Passing
    def _forward_messages(self, observations, parameters, forward_message,
            weights=None, tqdm=None, only_return_last=False):
        # Return list of forward messages Pr(x_{t} | y_{<=t})
        # y is num_obs x m matrix
        num_obs = np.shape(observations)[0]
        if not only_return_last:
            forward_messages = [None]*(num_obs+1)
            forward_messages[0] = forward_message

        mean_precision = forward_message['mean_precision']
        precision = forward_message['precision']
        log_constant = forward_message['log_constant']

        A = parameters.A
        LQinv = parameters.LQinv
        Qinv = np.dot(LQinv, LQinv.T)
        AtQinv = np.dot(A.T, Qinv)
        AtQinvA = np.dot(AtQinv, A)
        C = parameters.C
        LRinv = parameters.LRinv
        Rinv = np.dot(LRinv, LRinv.T)
        CtRinv = np.dot(C.T, Rinv)
        CtRinvC = np.dot(CtRinv, C)

        pbar = range(num_obs)
        if tqdm is not None:
            pbar = tqdm(pbar, leave=False)
            pbar.set_description("forward messages")
        for t in pbar:
            y_cur = observations[t]
            weight_t = 1.0 if weights is None else weights[t]

            # Calculate Predict Parameters
            J = np.linalg.solve(AtQinvA + precision, AtQinv)
            pred_mean_precision = np.dot(J.T, mean_precision)
            pred_precision = Qinv - np.dot(AtQinv.T, J)

            # Calculate Observation Parameters
            y_mean = np.dot(C,
                    np.linalg.solve(pred_precision, pred_mean_precision))
            y_precision = Rinv - np.dot(CtRinv.T,
                    np.linalg.solve(CtRinvC + pred_precision, CtRinv))
            log_c = (-0.5 * np.dot(y_cur-y_mean,
                            np.dot(y_precision, y_cur-y_mean)) + \
                     0.5 * np.linalg.slogdet(y_precision)[1] + \
                    -0.5 * self.m * np.log(2*np.pi))
            log_constant += log_c * weight_t

            # Calculate Filtered Parameters
            new_mean_precision = pred_mean_precision + np.dot(CtRinv, y_cur)
            new_precision = pred_precision + CtRinvC

            # Save Messages
            mean_precision = new_mean_precision
            precision = new_precision
            if not only_return_last:
                forward_messages[t+1] = {
                    'mean_precision': mean_precision,
                    'precision': precision,
                    'log_constant': log_constant,
                }
        if only_return_last:
            last_message = {
                    'mean_precision': mean_precision,
                    'precision': precision,
                    'log_constant': log_constant,
                }
            return last_message
        else:
            return forward_messages

    def _backward_messages(self, observations, parameters, backward_message,
            weights=None, tqdm=None, only_return_last=False):
        # Return list of backward messages Pr(y_{>t} | x_t)
        # y is num_obs x n matrix
        num_obs = np.shape(observations)[0]
        if not only_return_last:
            backward_messages = [None]*(num_obs+1)
            backward_messages[-1] = backward_message

        mean_precision = backward_message['mean_precision']
        precision = backward_message['precision']
        log_constant = backward_message['log_constant']

        A = parameters.A
        LQinv = parameters.LQinv
        Qinv = np.dot(LQinv, LQinv.T)
        AtQinv = np.dot(A.T, Qinv)
        AtQinvA = np.dot(AtQinv, A)
        C = parameters.C
        LRinv = parameters.LRinv
        Rinv = np.dot(LRinv, LRinv.T)
        CtRinv = np.dot(C.T, Rinv)
        CtRinvC = np.dot(CtRinv, C)

        pbar = reversed(range(num_obs))
        if tqdm is not None:
            pbar = tqdm(pbar, total=num_obs, leave=False)
            pbar.set_description("backward messages")
        for t in pbar:
            y_cur = observations[t]
            weight_t = 1.0 if weights is None else weights[t]

            # Helper Values
            xi = Qinv + precision + CtRinvC
            L = np.linalg.solve(xi, AtQinv.T)
            vi = mean_precision + np.dot(CtRinv, y_cur)

            # Calculate new parameters
            log_c = (-0.5 * self.m * np.log(2.0*np.pi) + \
                    np.sum(np.log(np.diag(LRinv))) + \
                    np.sum(np.log(np.diag(LQinv))) + \
                    -0.5 * np.linalg.slogdet(xi)[1] + \
                    -0.5 * np.dot(y_cur, np.dot(Rinv, y_cur)) + \
                    0.5 * np.dot(vi, np.linalg.solve(xi, vi)))

            log_constant += log_c * weight_t

            new_mean_precision = np.dot(L.T, vi)
            new_precision = AtQinvA - np.dot(AtQinv, L)

            # Save Messages
            mean_precision = new_mean_precision
            precision = new_precision

            if not only_return_last:
                backward_messages[t] = {
                    'mean_precision': mean_precision,
                    'precision': precision,
                    'log_constant': log_constant,
                }
        if only_return_last:
            last_message = {
                    'mean_precision': mean_precision,
                    'precision': precision,
                    'log_constant': log_constant,
                }
            return last_message
        else:
            return backward_messages

    ## Loglikelihood Functions
    def marginal_loglikelihood(self, observations, parameters,
            forward_message=None, backward_message=None, weights=None,
            tqdm=None, **kwargs):
        # Run forward pass + combine with backward pass
        # y is num_obs x n matrix
        if forward_message is None:
            forward_message = self.default_forward_message
        if backward_message is None:
            backward_message = self.default_backward_message

        # forward_pass is Pr(x_{T-1} | y_{<=T-1})
        forward_message = self._forward_message(
                observations=observations,
                parameters=parameters,
                forward_message=forward_message,
                weights=weights,
                tqdm=tqdm,
                **kwargs)

        # Calculate the marginal loglikelihood of forward + backward message
        f_mean_precision = forward_message['mean_precision']
        f_precision = forward_message['precision']
        c_mean_precision = f_mean_precision + backward_message['mean_precision']
        c_precision = f_precision + backward_message['precision']
        weight = 1.0 if weights is None else weights[-1]

        loglikelihood = forward_message['log_constant'] + \
                (backward_message['log_constant'] + \
                +0.5 * np.linalg.slogdet(f_precision)[1] + \
                -0.5 * np.linalg.slogdet(c_precision)[1] + \
                -0.5 * np.dot(f_mean_precision,
                        np.linalg.solve(f_precision, f_mean_precision)
                    ) + \
                0.5 * np.dot(c_mean_precision,
                    np.linalg.solve(c_precision, c_mean_precision)
                    )
                ) * weight

        return loglikelihood

    def complete_data_loglikelihood(self, observations, latent_vars, parameters,
            forward_message=None, weights=None, **kwargs):
        if forward_message is None:
            forward_message = self.default_forward_message

        log_constant = 0.0
        A = parameters.A
        LQinv = parameters.LQinv
        C = parameters.C
        LRinv = parameters.LRinv

        x_prev = forward_message.get('x_prev')
        for t, (y_t, x_t) in enumerate(zip(observations, latent_vars)):
            weight_t = 1.0 if weights is None else weights[t]

            # Pr(X_t | X_t-1)
            if (x_prev is not None):
                diffLQinv = np.dot(x_t - np.dot(A,x_prev), LQinv)
                log_c = (-0.5 * self.n * np.log(2*np.pi) + \
                        -0.5 * np.dot(diffLQinv, diffLQinv) + \
                        np.sum(np.log(np.diag(LQinv))))
                log_constant += log_c * weight_t

            # Pr(Y_t | X_t)
            LRinvTymCx = np.dot(LRinv.T, y_t - np.dot(C, x_t))
            log_c = (-0.5 * self.m * np.log(2*np.pi) + \
                    -0.5*np.dot(LRinvTymCx, LRinvTymCx) + \
                    np.sum(np.log(np.diag(LRinv))))
            log_constant += log_c * weight_t
            x_prev = x_t

        return log_constant

    def predictive_loglikelihood(self, observations, parameters, lag=1,
            forward_message=None, backward_message=None, tqdm=None, **kwargs):
        if forward_message is None:
            forward_message = self.default_forward_message
        if backward_message is None:
            backward_message = self.default_backward_message

        # Calculate Filtered
        if lag == 0:
            forward_messages = self.forward_pass(observations,
                    parameters, forward_message, tqdm=tqdm, **kwargs)
        else:
            forward_messages = self.forward_pass(observations[0:-lag],
                    parameters, forward_message, tqdm=tqdm, **kwargs)
        loglike = 0.0
        A = parameters.A
        Q = parameters.Q
        C = parameters.C
        R = parameters.R
        pbar = range(lag, np.shape(observations)[0])
        if tqdm is not None:
            pbar = tqdm(pbar, leave=False)
            pbar.set_description('predictive_loglikelihood')
        for t in pbar:
            # Calculate Pr(x_t | y_{<=t-lag}, theta)
            mean_precision = forward_messages[t-lag]['mean_precision']
            precision = forward_messages[t-lag]['precision']
            mean = np.linalg.solve(precision, mean_precision)
            var = np.linalg.inv(precision)
            for l in range(lag):
                mean = np.dot(A, mean)
                var = np.dot(A, np.dot(var, A.T)) + Q

            y_mean = np.dot(C, mean)
            y_var = np.dot(C, np.dot(var, C.T)) + R
            y_cur = observations[t]
            log_like_t = -0.5 * np.dot(y_cur - y_mean,
                    np.linalg.solve(y_var, y_cur - y_mean)) + \
                        -0.5 * np.linalg.slogdet(y_var)[1] + \
                        -0.5 * self.m * np.log(2*np.pi)
            loglike += log_like_t
        return loglike

    ## Gradient Functions
    def gradient_marginal_loglikelihood(self, observations, parameters,
            forward_message=None, backward_message=None, weights=None,
            include_init=True, tqdm=None):
        A, LQinv, C, LRinv = \
                parameters.A, parameters.LQinv, parameters.C, parameters.LRinv

        # Forward Pass
        # forward_messages = [Pr(x_{t} | y_{-inf:t}), y{t}] for t = -1,...,T-1
        forward_messages = self.forward_pass(observations,
                parameters, forward_message,
                include_init_message=True)

        # Backward Pass
        # backward_messages = [Pr(y_{t+1:inf} | x_{t}), Y_{t}] for t =-1,...,T-1
        backward_messages = self.backward_pass(observations,
                parameters, backward_message,
                include_init_message=True)

        # Gradients
        A_grad = np.zeros_like(A)
        C_grad = np.zeros_like(C)
        LQinv_grad = np.zeros_like(LQinv)
        LRinv_grad = np.zeros_like(LRinv)

        # Helper Constants
        Rinv = np.dot(LRinv, LRinv.T)
        RinvC = np.dot(Rinv, C)
        CtRinvC = np.dot(C.T, RinvC)
        LRinv_diaginv = np.diag(np.diag(LRinv)**-1)

        Qinv = np.dot(LQinv, LQinv.T)
        QinvA = np.dot(Qinv, A)
        AtQinvA = np.dot(A.T, QinvA)
        LQinv_diaginv = np.diag(np.diag(LQinv)**-1)

        # Emission Gradients
        p_bar = zip(forward_messages[1:], backward_messages[1:], observations)
        if tqdm is not None:
            p_bar = tqdm(p_bar, total=np.shape(observations)[0], leave=False)
            p_bar.set_description("gradient loglike")
        for t, (forward_t, backward_t, y_t) in enumerate(p_bar):
            weight_t = 1.0 if weights is None else weights[t]

            # Pr(x_t | y)
            c_mean_precision = \
                    forward_t['mean_precision'] + backward_t['mean_precision']
            c_precision = \
                    forward_t['precision'] + backward_t['precision']

            x_mean = np.linalg.solve(c_precision, c_mean_precision)
            xxt_mean = np.linalg.inv(c_precision) + np.outer(x_mean, x_mean)

            # Gradient of C
            C_grad += weight_t * (np.outer(np.dot(Rinv, y_t), x_mean) + \
                    -1.0 * np.dot(RinvC, xxt_mean))

            # Gradient of LRinv
            Cxyt = np.outer(np.dot(C, x_mean), y_t)
            CxxtCt = np.dot(C, np.dot(xxt_mean, C.T))
            LRinv_grad += weight_t * (LRinv_diaginv + \
                -1.0*np.dot(np.outer(y_t, y_t) - Cxyt - Cxyt.T + CxxtCt, LRinv)
                )

        # Transition Gradients
        if include_init:
            pbar = zip(
                forward_messages[:-1], backward_messages[1:], observations)
        else:
            pbar = zip(
                forward_messages[1:-1], backward_messages[2:], observations[1:])
        for t, (forward_t, backward_t, y_t) in enumerate(pbar):
            weight_t = 1.0 if weights is None else weights[t]
            # Pr(x_t, x_t+1 | y)
            c_mean_precision = \
                np.concatenate([
                    forward_t['mean_precision'],
                    backward_t['mean_precision'] + np.dot(RinvC.T,y_t)
                    ])
            c_precision = \
                np.block([
                    [forward_t['precision'] + AtQinvA, -QinvA.T],
                    [-QinvA, backward_t['precision'] + CtRinvC + Qinv]
                    ])

            c_mean = np.linalg.solve(c_precision, c_mean_precision)
            c_cov = np.linalg.inv(c_precision)

            xp_mean = c_mean[0:self.n]
            xn_mean = c_mean[self.n:]
            xpxpt_mean = c_cov[0:self.n, 0:self.n] + np.outer(xp_mean, xp_mean)
            xnxpt_mean = c_cov[self.n:, 0:self.n] + np.outer(xn_mean, xp_mean)
            xnxnt_mean = c_cov[self.n:, self.n:] + np.outer(xn_mean, xn_mean)

            # Gradient of A
            A_grad += weight_t * np.dot(Qinv, xnxpt_mean - np.dot(A,xpxpt_mean))

            # Gradient of LQinv
            Axpxnt = np.dot(A, xnxpt_mean.T)
            AxpxptAt = np.dot(A, np.dot(xpxpt_mean, A.T))
            LQinv_grad += weight_t * (LQinv_diaginv + \
                -1.0*np.dot(xnxnt_mean - Axpxnt - Axpxnt.T + AxpxptAt, LQinv))

        grad = dict(
                A=A_grad,
                LQinv_vec=LQinv_grad[np.tril_indices_from(LQinv_grad)],
                C=C_grad,
                LRinv_vec=LRinv_grad[np.tril_indices_from(LRinv_grad)],
                )
        return grad

    def gradient_complete_data_loglikelihood(self, observations, latent_vars,
            parameters, forward_message=None, weights=None, tqdm=None,
            **kwargs):
        if forward_message is None:
            forward_message = self.default_forward_message
        A = parameters.A
        LQinv = parameters.LQinv
        Qinv = parameters.Qinv
        LQinv_Tinv = lower_tri_mat_inv(LQinv).T
        C = parameters.C
        LRinv = parameters.LRinv
        Rinv = parameters.Rinv
        LRinv_Tinv = lower_tri_mat_inv(LRinv).T

        # Gradients
        grad = {var: np.zeros_like(value)
                for var, value in parameters.as_dict().items()}

        if len(np.shape(latent_vars)) == 2:
            # Only One Sample
            # Transition Gradients
            x_prev = forward_message.get('x_prev')
            for t, x_t in enumerate(latent_vars):
                weight_t = 1.0 if weights is None else weights[t]
                if x_prev is not None:
                    diff = x_t - np.dot(A, x_prev)
                    grad['A'] += weight_t * np.outer(np.dot(Qinv, diff), x_prev)
                    grad['LQinv'] += weight_t * (LQinv_Tinv + \
                        -1.0*np.dot(np.outer(diff, diff), LQinv))
                x_prev = x_t

            # Emission Gradients
            for t, (x_t, y_t) in enumerate(zip(latent_vars, observations)):
                weight_t = 1.0 if weights is None else weights[t]
                diff = y_t - np.dot(C, x_t)
                grad['C'] += weight_t * np.outer(np.dot(Rinv, diff), x_t)
                grad['LRinv'] += weight_t * (LRinv_Tinv + \
                    -1.0*np.dot(np.outer(diff, diff), LRinv))
        elif len(np.shape(latent_vars)) == 3:
            # Average over Multiple Latent Vars
            num_samples = np.shape(latent_vars)[2]

            # Transition Gradients
            x_prev = forward_message.get('x_prev')
            for t, x_t in enumerate(latent_vars):
                weight_t = 1.0 if weights is None else weights[t]
                if x_prev is not None:
                    diff = x_t - np.dot(A, x_prev)
                    grad['A'] += weight_t * np.dot(Qinv,
                            np.dot(diff, x_prev.T))/num_samples
                    grad['LQinv'] += weight_t * (LQinv_Tinv + \
                        -1.0*np.dot(np.dot(diff, diff.T), LQinv)/num_samples)
                x_prev = x_t

            # Emission Gradients
            for t, (x_t, y_t_) in enumerate(zip(latent_vars, observations)):
                y_t = np.array([y_t_ for _ in range(num_samples)]).T
                weight_t = 1.0 if weights is None else weights[t]

                diff = y_t - np.dot(C, x_t)
                grad['C'] += weight_t * (np.dot(Rinv,
                    np.dot(diff, x_t.T))/num_samples)
                grad['LRinv'] += weight_t * (LRinv_Tinv + \
                    -1.0*np.dot(np.dot(diff, diff.T), LRinv)/num_samples)
        else:
            raise ValueError("Incorrect latent_var shape")

        grad['LQinv_vec'] = grad.pop('LQinv')[np.tril_indices(self.n)]
        grad['LRinv_vec'] = grad.pop('LRinv')[np.tril_indices(self.m)]
        return grad

    def gradient_loglikelihood(self, kind='marginal', **kwargs):
        if kind == 'marginal':
            return self.gradient_marginal_loglikelihood(**kwargs)
        elif kind == 'complete':
            return self.gradient_complete_data_loglikelihood(**kwargs)
        else:
            raise ValueError("Unrecognized `kind' {0}".format(kind))

    ## Gibbs Functions
    def calc_gibbs_sufficient_statistic(self, observations, latent_vars,
            **kwargs):
        """ Gibbs Sample Sufficient Statistics
        Args:
            observations (ndarray): num_obs observations
            latent_vars (ndarray): latent vars

        Returns:
            sufficient_stat (dict)
        """
        x = latent_vars
        y = observations

        # Sufficient Statistics for A and Q
        # From Emily Fox's Thesis Page 147
        PsiT = x[1:]
        PsiT_prev = x[:-1]
        transition_count = len(PsiT)
        Sx_prevprev = PsiT_prev.T.dot(PsiT_prev)
        Sx_curprev = PsiT.T.dot(PsiT_prev)
        Sx_curcur = PsiT.T.dot(PsiT)

        # Sufficient Statistics for C and R
        # From Emily Fox's Thesis Page 147
        PsiT = y
        PsiT_prev = x
        emission_count = len(PsiT)
        S_prevprev = PsiT_prev.T.dot(PsiT_prev)
        S_curprev = PsiT.T.dot(PsiT_prev)
        S_curcur = PsiT.T.dot(PsiT)

        # Return sufficient Statistics
        sufficient_stat = {}
        sufficient_stat['A'] = dict(
                S_prevprev=Sx_prevprev,
                S_curprev=Sx_curprev,
                )
        sufficient_stat['Q'] = dict(
                S_count=transition_count,
                S_prevprev=Sx_prevprev,
                S_curprev=Sx_curprev,
                S_curcur=Sx_curcur,
                )
        sufficient_stat['R'] = dict(
                S_count=emission_count,
                S_prevprev=S_prevprev,
                S_curprev=S_curprev,
                S_curcur=S_curcur,
                )
        sufficient_stat['C'] = dict(
                S_prevprev=S_prevprev,
                S_curprev=S_curprev,
                )
        return sufficient_stat

    ## Predict Functions
    def latent_var_distr(self, observations, parameters,
            distr='marginal', lag=None,
            forward_message=None, backward_message=None,
            tqdm=None):
        if distr != 'marginal':
            raise NotImplementedError()
        if forward_message is None:
            forward_message = self.default_forward_message
        if backward_message is None:
            backward_message = self.default_backward_message

        # Setup Output
        L = np.shape(observations)[0]
        mean = np.zeros((L, self.n))
        cov = np.zeros((L, self.n, self.n))

        # Forward Pass
        forward_messages = self.forward_pass(
                observations=observations,
                parameters=parameters,
                forward_message=forward_message,
                tqdm=tqdm
                )

        pbar = range(L)
        if tqdm is not None:
            pbar = tqdm(pbar, leave=False)
            pbar.set_description('calc latent var distr')

        if lag is None:
            # Smoothing
            backward_messages = self.backward_pass(
                observations=observations,
                parameters=parameters,
                backward_message=backward_message,
                tqdm=tqdm
                )
            for t in pbar:
                mean_precision = \
                        forward_messages[t]['mean_precision'] + \
                        backward_messages[t]['mean_precision']
                precision = \
                        forward_messages[t]['precision'] + \
                        backward_messages[t]['precision']

                mean[t] = np.linalg.solve(precision, mean_precision)
                cov[t] = np.linalg.inv(precision)
            return mean, cov

        elif lag <= 0:
            # Prediction/Filtering
            A, Q = parameters.A, parameters.Q
            for t in pbar:
                if t+lag >= 0:
                    mean_precision = forward_messages[t+lag]['mean_precision']
                    precision = forward_messages[t+lag]['precision']
                else:
                    mean_precision = forward_message['mean_precision']
                    precision = forward_message['precision']
                mean_lag = np.linalg.solve(precision, mean_precision)
                cov_lag = np.linalg.inv(precision)

                # Forward Simulate
                for _ in range(-lag):
                    mean_lag = np.dot(A, mean_lag)
                    cov_lag = np.dot(np.dot(A, cov_lag), A.T) + Q

                mean[t] = mean_lag
                cov[t] = cov_lag
            return mean, cov

        else:
            # Fixed-lag Smoothing
            for t in pbar:
                # Backward Messages
                back_obs = observations[t:min(t+lag, L)]
                fixed_lag_message = self.backward_message(
                        observations=back_obs,
                        parameters=parameters,
                        backward_message=backward_message,
                        )
                # Output
                mean_precision = \
                        forward_messages[t]['mean_precision'] + \
                        fixed_lag_message['mean_precision']
                precision = \
                        forward_messages[t]['precision'] + \
                        fixed_lag_message['precision']
                mean[t] = np.linalg.solve(precision, mean_precision)
                cov[t] = np.linalg.inv(precision)
            return mean, cov

    def latent_var_sample(self, observations, parameters,
            forward_message=None, backward_message=None,
            distr='joint', lag=None, num_samples=None,
            tqdm=None, include_init=False, **kwargs):
        if forward_message is None:
            forward_message = self.default_forward_message
        if backward_message is None:
            backward_message = self.default_backward_message
        if distr == 'joint' and lag is not None:
            raise ValueError("Must set distr to 'marginal' for lag != None")

        A, LQinv = parameters.A, parameters.LQinv
        AtQinv = np.dot(A.T, np.dot(LQinv, LQinv.T))
        AtQinvA = np.dot(AtQinv, A)

        if distr == 'joint':
            # Forward Pass
            forward_messages = self.forward_pass(
                    observations=observations,
                    parameters=parameters,
                    forward_message=forward_message,
                    include_init_message=include_init,
                    tqdm=tqdm
                    )
            L = len(forward_messages)

            if num_samples is not None:
                x = np.zeros((L, self.n, num_samples))
            else:
                x = np.zeros((L, self.n))

            # Backward Sampler
            x_cov = np.linalg.inv(forward_messages[-1]['precision'])
            x_mean = np.dot(x_cov, forward_messages[-1]['mean_precision'])
            x[-1] = np.random.multivariate_normal(mean=x_mean, cov=x_cov,
                    size=num_samples).T

            pbar = reversed(range(L-1))
            if tqdm is not None:
                pbar = tqdm(pbar, leave=False)
                pbar.set_description("backward smoothed sampling x")
            for t in pbar:
                x_next = x[t+1]
                x_cov = np.linalg.inv(forward_messages[t]['precision'] +
                        AtQinvA)
                if num_samples is None:
                    x_mean = np.dot(x_cov,
                            forward_messages[t]['mean_precision'] + \
                                np.dot(AtQinv, x_next))
                else:
                    x_mean = np.dot(x_cov,
                            np.outer(forward_messages[t]['mean_precision'],
                                np.ones(num_samples)) + \
                                np.dot(AtQinv, x_next))
                x[t] = x_mean + np.random.multivariate_normal(
                        mean=np.zeros(self.n), cov=x_cov, size=num_samples,
                        ).T
            return x

        elif distr == 'marginal':
            # Calculate Distribution
            x_mean, x_cov = self.latent_var_distr(observations, parameters,
                    lag=lag, forward_message=forward_message,
                    backward_message=backward_message, tqdm=tqdm,
                    )
            # Sample from Distribution
            L = x_mean.shape[0]
            if num_samples is not None:
                x = np.zeros((x_mean.shape[0], self.n, num_samples))
            else:
                x = np.zeros((x_mean.shape[0], self.n))
            pbar = reversed(range(L))
            if tqdm is not None:
                pbar = tqdm(pbar, leave=False)
                pbar.set_description("sampling x")
            for t in pbar:
                x[t] = x_mean[t] + np.random.multivariate_normal(
                        mean=np.zeros(self.n), cov=x_cov[t], size=num_samples,
                        ).T

        else:
            raise ValueError("Unrecognized `distr'; {0}".format(distr))
        return

    def latent_var_pairwise_marginal(self, observations, parameters,
            forward_message=None, backward_message=None,
            distribution='smoothed', tqdm=None):
        """ Calculate latent var marginal distribution

        Args:
            observations (ndarray): num_obs by n observations
            parameters (LGSSMParameters): parameters
            forward_message (dict): alpha message
                (e.g. Pr(x_{-1} | y_{-inf:-1}))
            backward_message (dict): beta message
                (e.g. Pr(y_{T:inf} | x_{T-1}))
                'likelihood_vector' (ndarray) dimension num_states
                'y_next' (ndarray) dimension p by m, optional
            distr (string): 'smoothed', 'filtered', 'predict'
                smoothed: Pr(X | Y, theta)
                filtered: Pr(X_t | Y_<=t, theta)
                predictive:Pr(X_t | Y_<t, theta)

        Returns
            mean (ndarray): num_obs by 2n, pairwise mean (x_t, x_{t+1})
            cov (ndarray): num_obs by 2n by  2n, pairwise covariance
        """
        if forward_message is None:
            forward_message = self.default_forward_message
        if backward_message is None:
            backward_message = self.default_backward_message

        num_obs = np.shape(observations)[0]
        mean = np.zeros((num_obs, 2*self.n))
        cov = np.zeros((num_obs, 2*self.n, 2*self.n))
        if distribution == 'smoothed':
            forward_messages = self.forward_pass(
                    observations=observations,
                    parameters=parameters,
                    forward_message=forward_message,
                    include_init_message=True,
                    tqdm=tqdm
                    )
            backward_messages = self.backward_pass(
                    observations=observations,
                    parameters=parameters,
                    backward_message=backward_message,
                    include_init_message=True,
                    tqdm=tqdm
                    )
            # Helper Constants
            C = parameters.C
            Rinv = parameters.Rinv
            RinvC = np.dot(Rinv, C)
            CtRinvC = np.dot(C.T, RinvC)

            A = parameters.A
            Qinv = parameters.Qinv
            QinvA = np.dot(Qinv, A)
            AtQinvA = np.dot(A.T, QinvA)

            for t in range(num_obs):
                y_t = observations[t]
                mean_precision = \
                    np.concatenate([
                        forward_messages[t]['mean_precision'],
                        backward_messages[t+1]['mean_precision'] + \
                                np.dot(RinvC.T, y_t)
                        ])
                precision = \
                    np.block([
                        [forward_messages[t]['precision'] + AtQinvA, -QinvA.T],
                        [-QinvA,
                        backward_messages[t+1]['precision'] + CtRinvC + Qinv]
                        ])

                mean[t] = np.linalg.solve(precision, mean_precision)
                cov[t] = np.linalg.inv(precision)
            return mean, cov

        elif distribution == 'filtered':
            raise NotImplementedError()
        elif distribution == 'predictive':
            raise NotImplementedError()

        else:
            raise ValueError("Invalid `distribution'; {0}".format(distribution))
        return

    def y_distr(self, observations, parameters,
            distr='marginal', lag=None,
            forward_message=None, backward_message=None,
            latent_vars=None, tqdm=None):
        if forward_message is None:
            forward_message = self.default_forward_message
        if backward_message is None:
            backward_message = self.default_backward_message
        if latent_vars is not None:
            x_mean, x_cov = self.latent_var_distr(
                    observations=observations,
                    parameters=parameters,
                    distr=distr,
                    lag=lag,
                    forward_message=forward_message,
                    backward_message=backward_message,
                    latent_vars=latent_vars,
                    tqdm=tqdm,
                    )
            C, R = parameters.C, parameters.R
            y_mean = x_mean.dot(C.T)
            y_cov = np.array([C.dot(x_cov_t).dot(C.T) + R for x_cov_t in x_cov])
            return y_mean, y_cov
        else:
            if lag is None or lag >= 0:
                C, R = parameters.C, parameters.R
                y_mean = latent_vars.dot(C.T)
                y_cov = np.array([R for _ in range(L)])
                return y_mean, y_cov
            else:
                L = observations.shape[0]
                y_mean = np.nan(L, self.m)
                y_cov = np.nan(L, self.m, self.m)
                A, Q = parameters.A, parameters.Q
                C, R = parameters.C, parameters.R

                # Apply the transition + noise for lag steps
                tran = np.eye(self.m)
                cov = np.zeros(self.m, self.m)
                for _ in range(-lag):
                    tran = A.dot(tran)
                    cov = A.dot(cov).dot(A.T) + Q
                tran = C.dot(tran)
                cov = C.dot(cov).dot(C.T) + R

                # Calculate mean + cov based on latent_var[t+lag]
                y_mean[-lag:] = latent_vars[:lag].dot(tran.T)
                y_cov[-lag:] = np.array([cov for _ in range(L+lag)])

                # Calculate mean + cov based on forward message
                default_mean = np.linalg.solve(forward_messages['precision'],
                        forward_message['mean_precision'])
                default_cov = np.linalg.inv(forward_message['precision'])
                for _ in range(-lag):
                    default_cov = A.dot(default_cov).dot(A.T) + Q
                default_cov = C.dot(default_cov).dot(C) + R
                y_mean[:-lag] = np.array([default_mean.dot(tran.T)])
                y_cov[:-lag] = np.array([default_cov for _ in range(-lag)])

                return y_mean, y_cov

    def y_sample(self, observations, parameters,
            distr='marginal', lag=None, num_samples=None,
            forward_message=None, backward_message=None,
            latent_var=None, tqdm=None):
        if forward_message is None:
            forward_message = self.default_forward_message
        if backward_message is None:
            backward_message = self.default_backward_message
        if latent_vars is not None:
            latent_vars = self.latent_var_sample(
                    observations=observations,
                    parameters=parameters,
                    distr=distr,
                    lag=lag,
                    num_samples=num_samples,
                    forward_message=forward_message,
                    backward_message=backward_message,
                    latent_vars=latent_vars,
                    tqdm=tqdm,
                    )

        if num_samples is not None:
            y = np.zeros((L, self.m, num_samples))
        else:
            y = np.zeros((L, self.m))

        C, R = parameters.C, parameters.R
        y = x.dot(C.T) + np.random.multivariate_normal(
            mean=np.zeros(self.m), cov=R, size=num_samples).T
        return y

    def simulate_distr(self, T, parameters, include_init=True,
            init_message=None, tqdm=None):
        if init_message is None:
            init_message = self.default_forward_message
        m, n = np.shape(parameters.C)
        A = parameters.A
        C = parameters.C
        Q = parameters.Q
        R = parameters.R

        # Outputs
        latent_vars_mean = np.zeros((T+1, n), dtype=float)
        latent_vars_cov = np.zeros((T+1, n, n), dtype=float)
        obs_mean = np.zeros((T+1, m), dtype=float)
        obs_cov = np.zeros((T+1, m, m), dtype=float)

        # Init
        latent_vars_mean[0] = np.linalg.solve(init_message['precision'],
                init_message['mean_precision'])
        latent_vars_cov[0] = np.linalg.inv(init_message['precision'])
        obs_mean[0] = np.dot(C, latent_vars_mean[0])
        obs_cov[0] = np.dot(C, latent_vars_cov[0]).dot(C.T) + R

        pbar = range(1,T+1)
        if tqdm is not None:
            pbar = tqdm(pbar, leave=False)
            pbar.set_description("simulating data")
        for t in pbar:
            latent_vars_mean[t] = np.dot(A, latent_vars_mean[t-1])
            latent_vars_cov[t] = np.dot(A, latent_vars_cov[t-1]).dot(A.T) + Q
            obs_mean[t] = np.dot(C, latent_vars_mean[t])
            obs_cov[t] = np.dot(C, latent_vars_cov[t]).dot(C.T) + R

        if include_init:
            return dict(
                obs_mean=obs_mean,
                obs_cov=obs_cov,
                latent_vars_mean=latent_vars_mean,
                latent_vars_cov=latent_vars_cov,
                )
        else:
            return dict(
                obs_mean=obs_mean[1:],
                obs_cov=obs_cov[1:],
                latent_vars_mean=latent_vars_mean[1:],
                latent_vars_cov=latent_vars_cov[1:],
                )

    def simulate(self, T, parameters, init_message=None, num_samples=None,
            include_init=True, tqdm=None):
        if init_message is None:
            init_message = self.default_forward_message
        m, n = np.shape(parameters.C)
        A = parameters.A
        C = parameters.C
        Q = parameters.Q
        R = parameters.R

        # Outputs
        if num_samples is not None:
            latent_vars = np.zeros((T+1, num_samples, n), dtype=float)
            obs_vars = np.zeros((T+1, num_samples, m), dtype=float)
        else:
            latent_vars = np.zeros((T+1, n), dtype=float)
            obs_vars = np.zeros((T+1, m), dtype=float)

        # Init
        latent_vars[0] = np.random.multivariate_normal(
                mean=np.linalg.solve(init_message['precision'],
                    init_message['mean_precision']),
                cov=np.linalg.inv(init_message['precision']),
                size=num_samples,
                )
        obs_vars[0] = np.dot(latent_vars[0], C.T) + \
            np.random.multivariate_normal(
                    mean=np.zeros(m), cov=R, size=num_samples)

        pbar = range(1,T+1)
        if tqdm is not None:
            pbar = tqdm(pbar, leave=False)
            pbar.set_description("simulating data")
        for t in pbar:
            latent_vars[t] = np.dot(latent_vars[t-1], A.T) + \
                np.random.multivariate_normal(
                        mean=np.zeros(n), cov=Q, size=num_samples)
            obs_vars[t] = np.dot(latent_vars[t], C.T) + \
                np.random.multivariate_normal(
                        mean=np.zeros(m), cov=R, size=num_samples)

        if num_samples is not None:
            obs_vars = np.swapaxes(obs_vars, 1, 2)
            latent_vars = np.swapaxes(latent_vars, 1, 2)

        if include_init:
            return dict(
                observations=obs_vars,
                latent_vars=latent_vars,
                )
        else:
            return dict(
                observations=obs_vars[1:],
                latent_vars=latent_vars[1:],
                )

    ## PF Functions
    def pf_loglikelihood_estimate(self, observations, parameters,
            subsequence_start=0, subsequence_end=None, weights=None,
            pf="poyiadjis_N", N=1000, kernel=None, forward_message=None,
            **kwargs):
        # Set kernel
        Kernel = self._get_kernel(kernel)

        # Prior Mean + Variance
        if forward_message is None:
            forward_message = self.default_forward_message
        prior_var = np.linalg.inv(forward_message['precision'])
        prior_mean = np.linalg.solve(prior_var,
                forward_message['mean_precision'])

        # Run buffered pf
        out = buffered_pf_wrapper(pf=pf,
                observations=observations,
                parameters=parameters,
                N=N,
                kernel=Kernel,
                additive_statistic_func=gaussian_sufficient_statistics,
                statistic_dim=self.n+2*self.n**2,
                t1=subsequence_start,
                tL=subsequence_end,
                weights=weights,
                prior_mean=prior_mean,
                prior_var=prior_var,
                **kwargs
                )
        loglikelihood = out['loglikelihood_estimate']
        return loglikelihood

    def pf_predictive_loglikelihood_estimate(self, observations, parameters,
            num_steps_ahead=1, subsequence_start=0, subsequence_end=None,
            weights=None,
            pf="filter", N=1000, kernel=None, forward_message=None,
            **kwargs):
        # Set kernel
        Kernel = self._get_kernel(kernel)

        # Prior Mean + Variance
        if forward_message is None:
            forward_message = self.default_forward_message
        prior_var = np.linalg.inv(forward_message['precision'])
        prior_mean = np.linalg.solve(prior_var,
                forward_message['mean_precision'])

        from functools import partial
        additive_statistic_func = partial(gaussian_predictive_loglikelihood,
                num_steps_ahead=num_steps_ahead,
                observations=observations,
                )

        # Run buffered pf
        out = buffered_pf_wrapper(pf=pf,
                observations=observations,
                parameters=parameters,
                N=N,
                kernel=Kernel,
                additive_statistic_func=additive_statistic_func,
                statistic_dim=num_steps_ahead+1,
                t1=subsequence_start,
                tL=subsequence_end,
                weights=weights,
                prior_mean=prior_mean,
                prior_var=prior_var,
                logsumexp=True,
                **kwargs
                )
        predictive_loglikelihood = out['statistics']
        predictive_loglikelihood[0] = out['loglikelihood_estimate']
        return predictive_loglikelihood

    def pf_gradient_estimate(self, observations, parameters,
            subsequence_start=0, subsequence_end=None, weights=None,
            pf="poyiadjis_N", N=1000, kernel=None, forward_message=None,
            **kwargs):
        # Set kernel
        Kernel = self._get_kernel(kernel)

        # Prior Mean + Variance
        if forward_message is None:
            forward_message = self.default_forward_message
        prior_var = np.linalg.inv(forward_message['precision'])
        prior_mean = np.linalg.solve(prior_var,
                forward_message['mean_precision'])

        # Run buffered pf
        complete_grad_dim = 2*self.n**2+self.n*self.m+self.m**2
        out = buffered_pf_wrapper(pf=pf,
                observations=observations,
                parameters=parameters,
                N=N,
                kernel=Kernel,
                additive_statistic_func=lgssm_complete_data_loglike_gradient,
                statistic_dim=complete_grad_dim,
                t1=subsequence_start,
                tL=subsequence_end,
                weights=weights,
                prior_mean=prior_mean,
                prior_var=prior_var,
                **kwargs
                )
        grad_estimate = average_statistic(out)
        if self.n*self.m > 1:
            grad = dict(
                LRinv_vec = grad_estimate[:(self.m+1)*self.m//2],
                LQinv_vec = grad_estimate[(self.m+1)*self.m//2:
                    (self.m+1)*self.m//2+(self.n+1)*self.n//2],
                C = np.reshape(grad_estimate[
                    (self.m+1)*self.m//2+(self.n+1)*self.n//2:
                    (self.m+1)*self.m//2+(self.n+1)*self.n//2+self.n*self.m],
                    (self.m, self.n),
                    ),
                A = np.reshape(grad_estimate[
                    (self.m+1)*self.m//2+(self.n+1)*self.n//2+self.n*self.m:],
                    (self.n, self.n),
                    ),
                )
        else:
            grad = dict(
                LRinv_vec = grad_estimate[0],
                LQinv_vec = grad_estimate[1],
                C = grad_estimate[2],
                A = grad_estimate[3],
                )
        return grad

    def pf_latent_var_distr(self, observations, parameters, lag=None,
            subsequence_start=0, subsequence_end=None, weights=None,
            pf="poyiadjis_N", N=1000, kernel=None, forward_message=None,
            **kwargs):
        if lag == 0 and pf != 'filter':
            raise ValueError("pf must be filter for lag = 0")
        elif lag is None and pf == 'filter':
            raise ValueError("pf must not be filter for smoothing")
        elif lag is not None and lag != 0:
            raise NotImplementedError("lag can only be None or 0")

        # Set kernel
        Kernel = self._get_kernel(kernel)

        # Prior Mean + Variance
        if forward_message is None:
            forward_message = self.default_forward_message
        prior_var = np.linalg.inv(forward_message['precision'])
        prior_mean = np.linalg.solve(prior_var,
                forward_message['mean_precision'])

        # Run buffered pf
        out = buffered_pf_wrapper(pf=pf,
                observations=observations,
                parameters=parameters,
                N=N,
                kernel=Kernel,
                additive_statistic_func=gaussian_sufficient_statistics,
                statistic_dim=self.n+2*self.n**2,
                t1=subsequence_start,
                tL=subsequence_end,
                weights=weights,
                prior_mean=prior_mean,
                prior_var=prior_var,
                elementwise_statistic=True,
                **kwargs
                )
        avg_statistic = average_statistic(out)
        if self.n > 1:
            avg_statistic = np.reshape(avg_statistic, (-1, self.n+2*self.n**2))
            x_mean = avg_statistic[:, 0:self.n]
            x_cov = np.reshape(
                avg_statistic[:, self.n:self.n+self.n**2],
                (-1, self.n, self.n),
                ) - np.einsum('ij,ik->ijk', x_mean, x_mean)
        else:
            avg_statistic = np.reshape(avg_statistic, (-1, 3))
            x_mean = avg_statistic[:, 0]
            x_cov = avg_statistic[:, 1] - x_mean**2

            x_mean = np.reshape(x_mean, (x_mean.shape[0], 1))
            x_cov = np.reshape(x_cov, (x_cov.shape[0], 1, 1))

        return x_mean, x_cov

    def _get_kernel(self, kernel):
        if kernel is None:
            if self.n*self.m == 1:
                kernel = 'optimal'
            else:
                kernel = 'highdim'
        if kernel == "prior":
            Kernel = LGSSMPriorKernel()
        elif kernel == "optimal":
            Kernel = LGSSMOptimalKernel()
        elif kernel == "highdim":
            Kernel = LGSSMHighDimOptimalKernel()
        else:
            raise ValueError("Unrecognized kernel = {0}".format(kernel))
        return Kernel

# Additive Statistics
def lgssm_complete_data_loglike_gradient(x_t, x_next, y_next, parameters,
        **kwargs):
    """ Gradient of Complete Data Log-Likelihood

    Gradient w/r.t. parameters of log Pr(y_{t+1}, x_{t+1} | x_t, parameters)

    Args:
        x_t (N by n ndarray): particles for x_t
        x_next (N by n ndarray): particles for x_{t+1}
        y_next (m ndarray): y_{t+1}
        parameters (Parameters): parameters
    Returns:
        grad_complete_data_loglike (N by p ndarray):
            gradient of complete data loglikelihood for particles
            [ grad_LRinv_vec, grad_LQinv_vec, grad_C, grad_A ]
    """
    N, n = np.shape(x_next)
    m = np.shape(y_next)[0]

    A = parameters.A
    LQinv = parameters.LQinv
    Qinv = parameters.Qinv
    C = parameters.C
    LRinv = parameters.LRinv
    Rinv = parameters.Rinv

    grad_complete_data_loglike = [None] * N
    if (n != 1) or (m != 1):
        LQinv_Tinv = np.linalg.inv(LQinv).T
        LRinv_Tinv = np.linalg.inv(LRinv).T
        for i in range(N):
            grad = {}
            diff = x_next[i] - np.dot(A, x_t[i])
            grad['A'] = np.outer(
                np.dot(Qinv, diff), x_t[i])
            grad['LQinv_vec'] = (
                    LQinv_Tinv + -1.0*np.dot(np.outer(diff, diff), LQinv)
                    )[np.tril_indices(n)]

            diff = y_next - np.dot(C, x_next[i])
            grad['C'] = np.outer(np.dot(Rinv, diff), x_next[i])
            grad['LRinv_vec'] = (
                    LRinv_Tinv + -1.0*np.dot(np.outer(diff, diff), LRinv)
                    )[np.tril_indices(m)]

            grad_complete_data_loglike[i] = np.concatenate([
                grad['LRinv_vec'].flatten(),
                grad['LQinv_vec'].flatten(),
                grad['C'].flatten(),
                grad['A'].flatten(),
                ])
        grad_complete_data_loglike = np.array(grad_complete_data_loglike)
    else:
        diff_x = x_next - A * x_t
        grad_A = Qinv * diff_x * x_t
        grad_LQinv = (LQinv**-1) - (diff_x**2) * LQinv
        diff_y = y_next - C * x_next
        grad_C = Rinv * diff_y * x_next
        grad_LRinv = (LRinv**-1) - (diff_y**2) * LRinv
        grad_complete_data_loglike = np.hstack([
            grad_LRinv, grad_LQinv, grad_C, grad_A])

    return grad_complete_data_loglike

def gaussian_predictive_loglikelihood(x_t, x_next, t, num_steps_ahead,
        parameters, observations,
        **kwargs):
    """ Predictive Log-Likelihood

    Calculate [Pr(y_{t+1+k} | x_{t+1} for k in [0,..., num_steps_ahead]]

    Args:
        x_t (N by n ndarray): particles for x_t
        x_next (N by n ndarray): particles for x_{t+1}
        num_steps_ahead
        parameters (Parameters): parameters
        observations (T by m ndarray): y
    Returns:
        predictive_loglikelihood (N by num_steps_ahead+1 ndarray)

    """
    N, n = np.shape(x_next)
    T, m = np.shape(observations)

    predictive_loglikelihood = np.zeros((N, num_steps_ahead+1))

    x_pred_mean = x_next + 0.0
    x_pred_cov = np.zeros((n, n))
    R, Q = parameters.R, parameters.Q
    for k in range(num_steps_ahead+1):
        if t+k >= T:
            break
        diff = (
            np.outer(np.ones(N), observations[t+k]) - \
            np.dot(x_pred_mean, parameters.C.T)
            )
        y_pred_cov = R + np.dot(parameters.C,
                np.dot(x_pred_cov, parameters.C.T))

        if m > 1:
            pred_loglike = (
                -0.5*np.sum(diff*np.linalg.solve(y_pred_cov, diff.T).T, axis=1)+\
                -0.5*m*np.log(2.0*np.pi) +\
                -0.5*np.linalg.slogdet(y_pred_cov)[1]
                    )
        else:
            pred_loglike = -0.5*diff**2/y_pred_cov + \
                    -0.5*np.log(2.0*np.pi) - 0.5*np.log(y_pred_cov)
            pred_loglike = pred_loglike[:,0]

        predictive_loglikelihood[:,k] = pred_loglike

        x_pred_mean = np.dot(x_pred_mean, parameters.A.T)
        x_pred_cov = Q + \
                np.dot(parameters.A,
                    np.dot(x_pred_cov, parameters.A.T))

    return predictive_loglikelihood

def gaussian_sufficient_statistics(x_t, x_next, y_next, **kwargs):
    """ Gaussian Sufficient Statistics

    h[0] = sum(x_{t+1})
    h[1] = sum(x_{t+1} x_{t+1}^T)
    h[2] = sum(x_t x_{t+1})

    Args:
        x_t (N by n ndarray): particles for x_t
        x_next (N by n ndarray): particles for x_{t+1}
        y_next (m ndarray): y_{t+1}
    Returns:
        h (N by p ndarray): sufficient statistic
    """
    N = np.shape(x_t)[0]
    if (len(np.shape(x_t)) > 1) and (np.shape(x_t)[1] > 1):
        # x is vector
        h = [x_next,
                np.einsum('ij,ik->ijk', x_next, x_next),
                np.einsum('ij,ik->ijk', x_t, x_next),
            ]
        h = np.hstack([np.reshape(h_, (N, -1)) for h_ in h])
    else:
        # n = 1, x is scalar
        h = np.hstack([x_next, x_next**2, x_t*x_next])
    return h


import plotly.express as px
import streamlit as st
import pandas as pd
import numpy as np
import element_generator as eg
import financials as fin
import time
import logging
from millify import millify

# Page Config
st.set_page_config(
    page_title="Equity Analysis",
    page_icon="ðŸ“ˆ",
    layout="wide",
    initial_sidebar_state="collapsed",
    menu_items={
        "About": 'Equity Analysis Tool'
    }
)

logging.basicConfig(level=logging.DEBUG, format='%(levelname)s - %(message)s')

input_container = st.container()  # Input (ticker)
info_container = st.container()   # Company Overview, Metrics, Candlestick Chart

def load_data(ticker):
    """
    Load data from the API and display it on the page.

    Parameters
    ----------
    ticker : str
        The ticker of the company to load data for.

    Returns
    -------
    None
    """
    logging.info('Loading data...')
    # Company Name, Address, Description
    overview, response_code_overview = fin.company_overview(ticker)
    # Daily Adjusted Prices
    daily_adjusted, response_code_daily = fin.daily_adjusted(ticker)

    if response_code_overview == 200 and response_code_daily == 200:
        try:
            info_container.empty()
            logging.info('Data loaded successfully.')
            # Company Info
            info_container.header(f"{overview['Name']} ({overview['Symbol']})")
            info_container.caption(f"{overview['Address']}")
            info_container.caption(f"{overview['Description']}")

            # --- Metrics --- #
            col1, col2, col3, col4 = info_container.columns(4)
            col1.metric(
                label='Price',
                value=f'USD {eg.get_latest_price(daily_adjusted)}'
            )
            col2.metric(
                label='Market Cap',
                value=f'USD {millify(overview["MarketCapitalization"], precision=3)}'
            )
            col3.metric(label='P/E Ratio', value=overview['PERatio'])
            col4.metric(label='Dividend Yield', value=overview['DividendYield'])

            # CandleStick Chart
            info_container.plotly_chart(
                eg.candlestick(daily_adjusted),
                use_container_width=True
            )

            # --- Income Statement & Balance Sheet --- #
            col1, col2 = info_container.columns(2)
            # Income Statement
            income_statement, response_code = fin.income_statement(ticker)
            if response_code == 200:
                col1.subheader('Income Statement')
                col1.dataframe(income_statement)

            # Balance Sheet
            balance_sheet, response_code = fin.balance_sheet(ticker)
            if response_code == 200:
                col2.subheader('Balance Sheet')
                col2.dataframe(balance_sheet)

        except Exception as e:
            # Print error message and stack trace
            logging.error(e, exc_info=True)
            info_container.error('API limit reached. Please try again in 13 seconds.')
    else:
        # NOTE: When exceeding API limits, the repsonse code is still 200,
        # but the response body contains an error message. Therefore, we
        # catch the error here and display it to the user.
        logging.error('Ticker not found. Please try again.')

# Input Container
with input_container:
    st.text_input('Enter a ticker', value='AAPL', max_chars=5, key='ticker')
    st.button('Submit', key='submit_button')

# Load data when the submit button is clicked
if st.session_state.submit_button:
    load_data(st.session_state.ticker)


import time
import argparse
import os
from PIL import Image
import pandas as pd
import numpy as np
import tensorflow as tf
from utils import setup_gpu_memory, get_image_names_with_path, join_images

parser = argparse.ArgumentParser(description='Process arguments.')
parser.add_argument('--data_dir', type=str, default='/projects/ncdot/2018/NC_2018_Images',
                    help='input dir of the videolog data to apply model prediction for')
parser.add_argument('--model_file', type=str,
                    default='/projects/ncdot/2018/machine_learning/model/guardrail_xception_subset.h5',
                    help='model file with path to be load for prediction')
parser.add_argument('--output_file', type=str,
                    default='/projects/ncdot/2018/machine_learning/output/guardrail_classification.csv',
                    help='prediction output csv file')
parser.add_argument('--image_base_names_file', type=str,
                    default='/projects/ncdot/2018/NC_2018_Images/metadata/sensor_data_mapped_unique_3584576.csv',
                    help='file that contains image base names of images in data_dir')
parser.add_argument('--image_base_name_length', type=int,
                    default='11',
                    help='the length of the image base name string')
parser.add_argument('--is_subset', action='store_true', default=False,
                    help='make predictions on subset data for development db ingestion')

args = parser.parse_args()
data_dir = args.data_dir
model_file = args.model_file
output_file = args.output_file
image_base_names_file = args.image_base_names_file
image_base_name_length = args.image_base_name_length
is_subset = args.is_subset

setup_gpu_memory()

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    # load the model
    model = tf.keras.models.load_model(model_file, compile=False)
print(model.summary())

def predict(image_base_name):
    path, left, front, right = get_image_names_with_path(data_dir, image_base_name)
    if not path or not left or not front or not right:
        return

    left_image = os.path.join(path, left)
    front_image = os.path.join(path, front)
    right_image = os.path.join(path, right)
    dst_img = join_images(left_image, front_image, right_image)
    if dst_img:
        img = dst_img.resize((299, 299), Image.ANTIALIAS)
        img = tf.keras.preprocessing.image.img_to_array(img)
        img = tf.keras.applications.xception.preprocess_input(img)
        predictions = model.predict(np.array([img]))
        return predictions[0][0]
    else:
        return np.nan

# read all image base names
if is_subset:
    df = pd.read_csv(image_base_names_file, dtype=str, header=0, index_col=False, nrows=18240,
                     usecols=["MAPPED_IMAGE"])
else:
    df = pd.read_csv(image_base_names_file, dtype=str, header=0, index_col=False, usecols=["MAPPED_IMAGE"])

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    ts = time.time()
    df['Probability'] = df.apply(lambda row: predict(row['MAPPED_IMAGE']), axis=1)
    te = time.time()
    print('time taken for model prediction:', te-ts)

df.to_csv(output_file, index=False)
print('Done')


#A template for the implementation of hclust
import math #sqrt
import numpy as np
import matplotlib.pyplot as plt
from scipy.cluster.hierarchy import dendrogram, linkage
import pandas as pd

#Importing Dataset
dataset = pd.read_csv('../data/movehubjoin.csv')

#Purchase Power
X = dataset.iloc[:,8]
#Cappuccino
Y = dataset.iloc[:,1]

#Puts all datapoints in points[] array
points = []
for i in range(sum(1 for line in open('../data/movehubjoin.csv'))):
    if (i < X.size) and (i < Y.size):
        point = [X[i], Y[i]]
        points.append(point)

#Purchase Power -- X-axis
X = dataset.iloc[:,8]
#Cappuccino 	   -- Y-axis
Y = dataset.iloc[:,1]

#Puts all datapoints in points[] array
#Each x- and y-coordinate stored in an array [x, y]
points = []
for i in range(sum(1 for line in open('../data/movehubjoin.csv'))):
	if (i < X.size) and (i < Y.size):
		point = [X[i], Y[i]]
		points.append(point)

# Accepts two data points a and b.
# Returns the distance between a and b.
# Note that this might be specific to your data.
def Distance(a,b):
    return math.sqrt( ((a[0]-b[0])**2)+((a[1]-b[1])**2) )

# Accepts two data points a and b.
# Produces a point that is the average of a and b.
def merge(a,b):
    return [(a[0] + b[0])/2 , (a[1] + b[1])/2]

# Accepts a list of data points.
# Returns the pair of points that are closest
def findClosestPair(D):
    if (len(D) <= 2):
        return D

    closestPair = [None] * 2
    minDistance = Distance(D[0], D[1])

    for dataPoint in D:
        c = D.index(dataPoint)
        for otherPoint in (D[:c]+D[c+1:]):
            if (Distance(dataPoint, otherPoint) < minDistance):
                minDistance = Distance(dataPoint, otherPoint)
                closestPair = [dataPoint, otherPoint]

    return closestPair

# Accepts a list of data points.
# Produces a tree structure corresponding to a
# Agglomerative Hierarchal clustering of D.
def HClust(D):
    centers = D
    splits = []

    while len(centers) > 1:
        # find closest pair [x, y] in D
        closestPair = findClosestPair(centers)

        #indexes i and (i+1)%2==1 in splits[] is the ith depth of dendogram
        #containing ith coordinates of closestPair
        #append closestPair according to ith depth of dendogram
        splits.append(closestPair[0])
        splits.append(closestPair[1])

        index = 0
        while index < len(centers):
            if centers[index] == closestPair[0]:
                #replace the element containing closestPair[x] in centers with the
                #mean point between the closestPair of points
                centers[index] = merge(closestPair[0], closestPair[1])

            elif centers[index] == closestPair[1]:
                #remove y-coordinate of closestPair from centers
                centers.remove(centers[index])

            index += 1

    dendoGram = visualize(splits)
    return dendoGram

#draws dendogram visualization of all points
def visualize(points):
    #turn regular array into numPy array
    numPyPoints = np.array(points)
    #DENDOGRAM
    linked = linkage(numPyPoints, 'single')
    plt.figure(figsize=(10, 7))
    dendrogram(linked,
                orientation='top',
                #labels=labelList,
                distance_sort='descending',
                show_leaf_counts=True)
    plt.xlabel('Purchase Power')
    plt.ylabel('Cappuccino')
    plt.savefig('../visualizations/hclust.png')
    return plt

def main():
    #create dendogram variable, then invoke matplotlib function show() on it
    myPlot = HClust(points)
    myPlot.show()

if __name__ == '__main__':
    main()


from typing import List

class Solution:
    def maxArea(self, height: List[int]) -> int:
        maxarea = 0
        left = 0
        right = len(height) - 1

        while left < right:
            width = right - left
            maxarea = max(maxarea, min(height[left], height[right]) * width)
            if height[left] <= height[right]:
                left += 1
            else:
                right -= 1
        return maxarea

if __name__ == '__main__':
    # Inputs
    height1 = [1,8,6,2,5,4,8,3,7]
    height2 = [1,1]

    solution1 = Solution()
    solution2 = Solution()

    test_1 = solution1.maxArea(height1)
    test_2 = solution1.maxArea(height2)

    print(f"Test 1: {test_1}")
    print(f"Test 2: {test_2}")


#!/usr/bin/env python
import roslib
roslib.load_manifest('learning_tf')
import rospy

import tf
import turtlesim.msg

def handle_turtle_pose(msg, turtlename): #callback function?
    br = tf.TransformBroadcaster()
    br.sendTransform((msg.x, msg.y, 0),
                     tf.transformations.quaternion_from_euler(0, 0, msg.theta),
                     rospy.Time.now(),
                     child= turtlename, # name of relative position of turtle
                     parent="world") # name of parent coordinate

if __name__ == '__main__':
    rospy.init_node('turtle_tf_broadcaster')
    turtlename = rospy.get_param('~turtle') # param from launch file
    test = rospy.get_param('~alo') # param from launch file
    print(test)
    rospy.Subscriber('/%s/pose' % turtlename, # get data from, for example, /turtle1/pose and /turtle2/pose. But may be turtle1, turtle2 do not created
                     turtlesim.msg.Pose,
                     handle_turtle_pose,
                     turtlename)
    rospy.spin()


#-*- coding:utf-8 -*-
import os
import json
import csv
from loguru import logger

class CsvWriter:

    #  Initialisierung
    def __init__(self, filename, fieldnames, broker):
        self.broker = broker
        self.broker.sub("iot-distance-sensor/data", self.write_line)
        self.filename = filename
        self.fieldnames = fieldnames
        # Need to os.path.exists before opening file -> opening file creates it if it is not there
        writeheaders = False
        if not os.path.exists(self.filename):
            logger.info("Creating the csv-file {} to store data", self.filename)
            writeheaders = True
        with open(self.filename, "a") as file:
            writer = csv.DictWriter(file, self.fieldnames)
            # Check if the csv-file already exists
            if writeheaders:
                writer.writeheader()

    #  line: ist die bereits formatierte Zeile, die nur noch geschrieben wird.
    def write_line (self, line):
        logger.trace("Writing line: {} to csv", line)
        with open(self.filename, "a") as file:
            writer = csv.DictWriter(file, self.fieldnames)
            writer.writerow(line)


import datetime

from internal.app.database.dao import user_dao, survey_dao, question_dao

class SurveyService:
    items_per_page = 20

    def get_survey_list_data(self, page):
        start = (page - 1) * self.items_per_page

        data = survey_dao.get_paginated_data(start, self.items_per_page)

        total_pages = len(data) // self.items_per_page + (
            1 if len(data) % self.items_per_page > 0 else 0
        )

        return data, total_pages

    @staticmethod
    def get_survey(sid):
        survey = survey_dao.get_one(sid)
        if not survey:
            return False
        return True

    @staticmethod
    def get_question(sid, page):
        question = question_dao.get_question_by_sid(sid, page)
        if not question:
            return False
        return question

    def check_answer(self, sid, page, answer):
        question = self.get_question(sid, page)
        return question.true_option == answer

    def create_survey(self, uid, data):
        data = {
            "user_id": uid,
            "title": data.get("title"),
            "created_at": datetime.datetime.now()
        }
        survey_dao.create(data)

        self.create_questions(data)

    @staticmethod
    def create_questions(data):
        questions = [key for key in data if key.startswith("question")]
        t_counter, f_counter = 1, 1
        for question in questions:
            data = {
                "body": data[question],
                "true_option": data[f"answer_option_{t_counter}"],
                "false_option_1": data[f"answer_option_{f_counter}"],
                "false_option_2": data[f"answer_option_{f_counter + 1}"],
                "false_option_3": data[f"answer_option_{f_counter + 2}"],
            }
            question_dao.create(data)
            t_counter += 1
            f_counter += 3


import psycopg2

from .config import config_db
from .queries import *

connection = None

class Database:

    def init_db(self):
        connection = psycopg2.connect(**config_db())

        return connection

    @classmethod
    def create_tables(cls):
        statements = [
            CREATE_TABLE_USERS,
            CREATE_TABLE_MEETUPS,
            CREATE_TABLE_RSVPS,
            CREATE_TABLE_QUESTIONS,
            CREATE_TABLE_TOKENS,
            CREATE_TABLE_QVOTES,
            CREATE_TABLE_COMMENTS]

        print("Creating Tables")

        for statement in statements:
            query_db(statement)
        print("Created Tables")

    @classmethod
    def drop_tables(cls):
        query_db(DROP_TABLES)

def query_db(statement, values=None, rowcount=False,
             return_value=False, one=False, many=False):
    result = None
    global connection

    try:
        connection = db_instance.init_db()
        cursor = connection.cursor()

        if one:
            cursor.execute(statement, values)
            result = cursor.fetchone()
        elif many:
            if values:
                cursor.execute(statement, values)
            else:
                cursor.execute(statement)
            result = cursor.fetchall()
        elif rowcount:
            cursor.execute(statement, values)
            result = cursor.rowcount
        else:
            cursor.execute(statement)

        connection.commit()
        cursor.close()

    except (Exception, psycopg2.DatabaseError) as er:
        print(er)
    finally:
        if (connection):
            connection.close()

    return result

db_instance = Database()


from gpglib.content_parsers.crypt import Mapped
from gpglib.content_parsers.base import Parser
from gpglib import errors

from Crypto.Hash import SHA

class CompressedParser(Parser):
    """Parse compressed packets"""

    def consume(self, tag, message, region):
        # Get the compression algorithm used
        # And the rest of the message
        algo, rest = region.readlist("uint:8, bytes")

        # Get decompressor and use it to decompress the packet
        decompressor = Mapped.compression.decompression[algo]
        uncompressed = decompressor(rest)

        # Parse the inner packet and return it
        return message.consume(uncompressed)

class SymEncryptedParser(Parser):
    """Parse symmetrically encrypted data packet"""

    def consume(self, tag, message, region, algo, session_key):
        # Get the encryption algorithm used
        cipher, _ = Mapped.algorithms.encryption[algo]

        # Find out the length of the IV
        iv_len = cipher.block_size + 2

        # Read in the encrypted IV
        # The ciphertext is what's left in `region` after the iv
        iv, ciphertext = region.readlist("bytes:%d, bytes" % iv_len)

        # Build the cipher object from the session key and the IV
        decryptor = cipher.new(session_key, cipher.MODE_OPENPGP, iv)

        # Decrypt the ciphertext
        decrypted = decryptor.decrypt(ciphertext)

        # Parse the inner packet and return it
        return message.consume(decrypted)

class LiteralParser(Parser):
    """Extracts various information from the packet but only returns the plaintext"""

    def consume(self, tag, message, region):
        # Is it binary ('b'), text ('t') or utf-8 text ('u')
        # The length of the filename in bytes
        format, filename_length = region.readlist("bytes:1, uint:8")

        # Read in filename
        # Read in the date (can mean anything. ie. creation, modification, or 0)
        filename, date = region.readlist("bytes:%d, uint:32" % filename_length)

        # Add the literal data to the list of decrypted plaintext
        message.add_plaintext(region.read("bytes"))

class UserIdParser(Parser):
    """Parses type-13 packets, which contain information about the user"""

    def consume(self, tag, message, region):
        message.userid = region.read("bytes")

class SymmetricallyEncryptedIntegrityProtectedDataPacketParser(Parser):
    """Parses type-18 packets"""

    def consume(self, tag, message, region, algo, session_key):
        version = region.read("uint:8")
        if version != 1:
            raise errors.PGPException("Expected integrity packet to have version 1")

        encrypted_data = region.read("bytes")

        # Get a cipher object we can use to decrypt the key (and fail if we can't)
        cipherkls, key_size = Mapped.algorithms.encryption[algo]

        block_size = cipherkls.block_size
        segment_size = block_size * 8
        cipher = cipherkls.new(
            session_key, cipherkls.MODE_CFB, b"\x00" * block_size, segment_size=segment_size
        )

        decrypted_iv = cipher.decrypt(encrypted_data[:block_size])
        decrypted_data = bytearray()

        offset = block_size
        first_block = cipher.decrypt(encrypted_data[offset : offset + block_size])

        offset += block_size
        iv_check = first_block[:2]
        if iv_check != decrypted_iv[-2:]:
            raise errors.PGPException("There was a mismatch in the integrity check")

        decrypted_data.extend(first_block[2:])

        padding = block_size - (len(encrypted_data) - offset) % block_size
        encrypted_data += b"\x00" * padding

        decrypted_data.extend(cipher.decrypt(encrypted_data[offset:]))
        decrypted_data = decrypted_data[:-padding]

        # We run into trouble here with partial body lengths and the MDC.
        # Let's make pretend.
        mdc_data = bytearray()
        main_data = bytes(decrypted_data)
        if decrypted_data[-22:-20] == b"\xd3\x14":
            # It's there. Hold it back.
            mdc_data = bytes(decrypted_data[-22:])
            main_data = bytes(decrypted_data[:-22])

        message.consume(main_data)
        if mdc_data:
            message.consume(mdc_data)

        if not hasattr(message, "mdc"):
            raise errors.PGPException(
                "Integrity protected message is missing modification detection code."
            )

        hash_ = SHA.new(decrypted_iv)
        hash_.update(decrypted_iv[-2:])
        hash_.update(decrypted_data[:-20])

        if message.mdc != hash_.digest():
            raise errors.PGPException(
                "Integrity protected message does not match modification detection code."
            )

class ModificationDetectionCodePacketParser(Parser):
    """Parses type-17 packets"""

    def consume(self, tag, message, region):
        # Just contains a 20 octet sha1 hash
        message.add_mdc(region.read("bytes:20"))

        # Hack to make sure this packet is last
        return dict(mdc_must_be_last=True)


import streamlit as st
def calculate_emi(p,n,r):
	emi=p*(r/100)*(1+(r/100)**n)/(((1+(r/100))**n)-1)
	st.write("EMI Calculated=",emi)

st.title("EMI Calculator")
principal=st.slider("Loan Amount",1000.0,1000000.0)
tenure=st.slider("Time Period of Loan",1,50)
roi=st.slider("Rate of Interest",1,20)
if st.button("Calculate"):
	calculate_emi(principal,tenure,roi)


import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split

# Leitura da base de dados
data = pd.read_csv('base_de_dados.csv')

# Definição das variáveis explicativas (X) e da variável resposta (y)
X = data[['idade', 'sexo', 'renda', 'historico_compras']]
y = data['contratou_x']

# Divisão da base de dados em treino e teste
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)

# Treinamento do modelo de regressão logística
model = LogisticRegression()
model.fit(X_train, y_train)

# Aplicação do modelo aos clientes que não contrataram o produto X
X_nao_contrataram = data[data['contratou_x'] == 0][['idade', 'sexo', 'renda', 'historico_compras']]
probabilidade_contratacao = model.predict_proba(X_nao_contrataram)

# Análise das diferenças de médias
media_gastos_contratantes = data[data['contratou_x'] == 1]['gastos'].mean()
media_gastos_nao_contratantes = data[data['contratou_x'] == 0]['gastos'].mean()

diferenca_medias = media_gastos_contratantes - media_gastos_nao_contratantes

print("Probabilidade de contratação do produto X para os clientes que não o contrataram:", probabilidade_contratacao[:, 1])
print("Diferença de médias de gastos entre os clientes que contrataram e os que não contrataram o produto X:", diferenca_medias)


import time
from pynput import mouse
from pynput import keyboard

M=mouse.Controller() #创建一个鼠标对象
K=keyboard.Controller() #创建一个键盘对象

print(M.position)

#找到QQ下拉框（1621，5）
M.position=(1400,0)

#找到对话联系人（1680，277）
time.sleep(0.5)
M.position=(1400,277)

#双击聊天框
M.click(mouse.Button.left, 2)

#定位到聊天栏（1059，625）
M.position=(800,625)
time.sleep(0.5)

#点击聊天栏
time.sleep(0.5)
M.click(mouse.Button.left, 1)
#输入对话
time.sleep(0.5)
K.type("hello,world!")
time.sleep(0.5)

#发送（ctrl+enter）
K.press(keyboard.Key.ctrl_l)
K.press(keyboard.Key.enter)
K.release(keyboard.Key.enter)
K.release(keyboard.Key.ctrl_l)


"""
we run here some mediapipe experiments.
"""
import argparse
from rt import env
import cv2
import time
import json
import mediapipe as mp
mp_face_detection = mp.solutions.face_detection
mp_drawing = mp.solutions.drawing_utils
def do_config(args):
    """print current configuration."""
    print(do_faces.__doc__)
    print(json.dumps(env.config,indent=2))

def do_faces(args):
    """face detection."""
    print(do_faces.__doc__)
    source=env.get_workpath("face-demographics-walking.mp4")
    cap = cv2.VideoCapture(source)
    # Check if camera opened successfully
    if (cap.isOpened()== False):
        print(f"Error opening video stream:{source}")
        return
    fps=float(cap.get(cv2.CAP_PROP_FPS) or 30)
    fps_delay=1/fps
    print(f"fps:{fps},fps_delay:{fps_delay}")
    with mp_face_detection.FaceDetection(
        model_selection=0,
        min_detection_confidence=0.5
        ) as face_detection:
        while True:
            msworked=time.time()
            ret, frame = cap.read()
            if ret==False:
                break
            #cv2.imshow(source, frame)
            # Flip the image horizontally for a later selfie-view display, and convert
            # the BGR image to RGB.
            image = cv2.cvtColor(cv2.flip(frame, 1), cv2.COLOR_BGR2RGB)
            # To improve performance, optionally mark the image as not writeable to
            # pass by reference.
            image.flags.writeable = False
            results = face_detection.process(image)

            # Draw the face detection annotations on the image.
            image.flags.writeable = True
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            if results.detections:
                for detection in results.detections:
                    mp_drawing.draw_detection(image, detection)
            cv2.imshow('MediaPipe Face Detection', image)
            msworked=time.time()-msworked
            msworked=int(msworked*1000)
            msdelay=max(1,int(fps_delay*1000)-msworked)
            #print("msworked:",msworked,msdelay)
            #time.sleep(fps_delay)
            if cv2.waitKey(msdelay) & 0xFF == ord('q'):
                break
    cap.release()
    cv2.destroyAllWindows()

if __name__=="__main__":
    jobs={k.replace("do_",""):v for k,v in locals().items() if k.startswith("do_")}
    ap=argparse.ArgumentParser(description=__doc__)
    ap.add_argument("-j","--job",required=True,help="job to perform", choices=jobs.keys())
    args=ap.parse_args()
    env.init()
    override={k:v for (k,v) in vars(args).items() if v is not None or k not in env.config}
    env.config.update(override)
    jobs[args.job](args)
    #with app.app_context():
    #    jobs[args.job]()


import argparse
import os
import shutil
import subprocess
import sys
from tempfile import mkdtemp

def showcase():
    assert os.path.exists(args.target), "Target folder %s doesn't exist" % args.target

    exts = ['.' + e for e in args.ext.split(",")]
    # source_files = [os.path.join(r, f) for r, fd, fl in os.walk(os.path.abspath(args.target)) for f in fl
    #                 if os.path.splitext(f)[1].lower() in exts]

    tmp = mkdtemp(prefix="showcase_")
    first = None
    c = 0
    for r, dirs, files in os.walk(os.path.abspath(args.target)):
        dirs.sort()
        for file in [os.path.join(r, f) for f in sorted(files) if os.path.splitext(f)[1].lower() in exts]:
            c += 1
            link_name = "%04i_%s" % (c, os.path.basename(file))
            if not first:
                first = os.path.join(tmp, link_name)
            os.symlink(file, os.path.join(tmp, link_name))
            print("%s -> %s" % (file, os.path.join(tmp, link_name)))

    subprocess.call(["xdg-open", first])

def permanent_showcase():
    SHOWCASE_NAME = "__showcase"
    target = args.target
    counter = {}

    def create_linker(_folder):
        if _folder not in counter:
            counter[_folder] = 0

        def _linker(item):
            if not args.skip_root or (_folder != target):
                c = counter[_folder]
                counter[_folder] += 1

                showcase_folder = os.path.join(_folder, SHOWCASE_NAME)
                os.makedirs(showcase_folder, exist_ok=True)
                link_name = os.path.join(showcase_folder, "%04d_%s" % (c, os.path.basename(item)))
                path_to_item = os.path.relpath(item, showcase_folder)
                if not os.path.exists(link_name):
                    os.symlink(path_to_item, link_name)

        return _linker

    def picture_list(folder, linkers):

        if args.rebuild or args.remove:
            sc_folder = os.path.join(folder, SHOWCASE_NAME)
            if os.path.exists(sc_folder):
                shutil.rmtree(sc_folder)

        for entry in sorted([e for e in os.scandir(folder) if e.name != SHOWCASE_NAME], key=lambda e: e.name):
            entry_path = os.path.join(folder, entry.name)
            if entry.is_dir():
                picture_list(entry_path, linkers + [create_linker(folder)])
            elif entry.is_file() and not args.remove:
                for linker in linkers:
                    linker(entry_path)

    picture_list(target, [print])

def catalog():
    source = os.path.expanduser(args.source)
    target = os.path.expanduser(args.target)
    assert os.path.exists(source), "Source folder doesn't exist"

    if not os.path.exists(target):
        os.makedirs(target, exist_ok=True)

    exts = ['.' + e for e in args.ext.split(",")]

    for root, folders, files in os.walk(source):
        print('scan %s...' % root)
        if os.path.basename(os.path.normpath(root))[0] == '.':
            continue

        relative_root = root[len(source):]
        for file in [f for f in files if os.path.splitext(f)[1].lower() in exts]:
            print('file %s...' % file)
            source_picture = os.path.join(root, file)
            target_picture = os.path.join(target, relative_root, file)

            target_exists = os.path.exists(target_picture)

            if not target_exists or args.force:
                if target_exists:
                    print('Replacing %s...' % target_picture)
                    os.remove(target_picture)
                print("Copying %s..." % os.path.join(relative_root, file), end='')
                sys.stdout.flush()
                os.makedirs(os.path.dirname(target_picture), exist_ok=True)
                shutil.copyfile(source_picture, target_picture)
                print("resizing...", end='')
                sys.stdout.flush()
                subprocess.call(["mogrify", "-resize", args.resize, target_picture])
                print('done')
            else:
                print("File %s already exist, skipped" % target_picture)

    print("done")

catalog_file = os.path.expanduser("~/.jpgutils-catalog.idx")

def _load_catalog_index():
    index = {}
    if os.path.exists(catalog_file):
        with open(catalog_file, 'r') as f:
            first = True
            for line in f.readlines():
                if first:
                    first = False
                    continue
                path, mtime = line[:-1].split('\t')
                index[path] = float(mtime)

    return index

def _write_catalog_index(index: dict):
    with open(catalog_file, 'w') as f:
        f.write("original_path\ttimestamp%n")

        for path, timestamp in index.items():
            f.write("%s\t%s\n" % (path, timestamp))

def catalog2():
    def path_mtime():
        return os.path.abspath(source_picture), os.path.getmtime(source_picture)

    def log(s, flush=False):
        if args.verbose:
            if flush:
                print(s, end='', flush=True)
            else:
                print(s)

    source = os.path.expanduser(args.source)
    target = os.path.expanduser(args.target)
    assert os.path.exists(source), "Source folder doesn't exist"

    index = _load_catalog_index()

    if not os.path.exists(target):
        os.makedirs(target, exist_ok=True)

    exts = ['.' + e for e in args.ext.split(",")]
    forced_resize = args.force

    sources = 0
    skipped = 0
    replaced = 0

    source = source[:-1] if source[-1] == '/' else source
    root_count = source.count('/')

    for root, folders, files in os.walk(source, topdown=True):
        folders.sort()
        level = root.count('/') - root_count
        _spacer = '  ' * level

        if os.path.basename(os.path.normpath(root))[0] == '.':
            continue

        if level > 0:
            log(_spacer + " - " + os.path.basename(root))

        target_root = os.path.join(target, root[len(source) + 1:])
        for file in sorted([f for f in files if os.path.splitext(f)[1].lower() in exts]):
            source_picture = os.path.join(root, file)
            target_picture = os.path.join(target_root, file)
            log('%s  -  %s -> ' % (_spacer, file), True)
            sources += 1

            path, mtime = path_mtime()
            source_changed = path in index and index[path] != mtime
            target_exists = os.path.exists(target_picture)

            if source_changed or forced_resize or not target_exists:
                repl = False
                if target_exists:
                    os.remove(target_picture)
                    replaced += 1
                    repl = True
                os.makedirs(os.path.dirname(target_picture), exist_ok=True)
                shutil.copyfile(source_picture, target_picture)
                subprocess.call(["mogrify", "-resize", args.resize, target_picture])
                if repl:
                    log("replaced.")
                else:
                    log('generated.')
            else:
                log("skipped.")
                skipped += 1

            index[os.path.abspath(source_picture)] = os.path.getmtime(source_picture)

    _write_catalog_index(index)
    print("""
Catalog generation done
    Source folder: %s
    Target folder: %s
    Found %d pictures
    Skipped %d pictures
    Generated %d pictures (%d replaced)
""" % (source, target, sources, skipped, (sources - skipped), replaced))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    subparsers = parser.add_subparsers()

    showcase_parser = subparsers.add_parser("showcase", help="Symlink all pictures in folder and open viewer")
    showcase_parser.set_defaults(command="showcase")
    showcase_parser.add_argument("target", help="File or directory to view")
    showcase_parser.add_argument("--ext", default="jpg,jpeg",
                                 help="Extensions valid to be shown, separated by comma")
    showcase_parser.add_argument("--permanent", help="Make showcase permanent")

    catalog_parser = subparsers.add_parser("catalog", help="Create catalog with smaller pictures")
    catalog_parser.set_defaults(command="catalog")
    catalog_parser.add_argument("source", help="Original directory from where creating the catalog")
    catalog_parser.add_argument("target", help="Final destination of the catalog")
    catalog_parser.add_argument("--resize", default="1920x1080",
                                help="Resize value of the pictures (default 1920x1080")
    catalog_parser.add_argument("--ext", default="jpg,jpeg",
                                help="Extensions valid to be copied, separated by comma")
    catalog_parser.add_argument("--force", action='store_true',
                                help='Force resized file to be overridden')

    catalog_parser2 = subparsers.add_parser("catalog2", help="Create catalog with smaller pictures")
    catalog_parser2.set_defaults(command="catalog2")
    catalog_parser2.add_argument("source", help="Original directory from where creating the catalog")
    catalog_parser2.add_argument("target", help="Final destination of the catalog")
    catalog_parser2.add_argument("--resize", default="1920x1080",
                                 help="Resize value of the pictures (default 1920x1080")
    catalog_parser2.add_argument("--ext", default="jpg,jpeg",
                                 help="Extensions valid to be copied, separated by comma")
    catalog_parser2.add_argument("--force", action='store_true',
                                 help='Force resized file to be overridden')
    catalog_parser2.add_argument('-v', "--verbose", action='store_true',
                                 help='Log all files encountered')

    pshowcase_parser = subparsers.add_parser("permanent-showcase", help="Create permanent showcase")
    pshowcase_parser.set_defaults(command="permanent_showcase")
    pshowcase_parser.add_argument("target", help="Final destination of the catalog")
    pshowcase_parser.add_argument("--skip-root", action="store_true", help="Skip root showcase folder")
    pshowcase_parser.add_argument("--rebuild", action="store_true", help="Delete all showcase folders and then create again")
    pshowcase_parser.add_argument("--remove", action="store_true", help="Delete all showcase folders")

    args = parser.parse_args()
    if args.command:
        if args.command == 'showcase':
            showcase()
        elif args.command == 'catalog':
            catalog()
        elif args.command == 'catalog2':
            catalog2()
        elif args.command == 'permanent_showcase':
            permanent_showcase()
        else:
            parser.print_usage()
    else:
        parser.print_usage()


from django.urls import path
from inec.views import (
    home,PollDetailView,local_gov_results,search,new_polling_unit
)
urlpatterns = [
    path('', home,name='home'),
    path('poll/unit/<int:pk>/',PollDetailView.as_view(),name='detail'),
    path('local/govt/result/',local_gov_results,name = 'results'),
    path('search/',search, name = 'search'),
    path('add/new/unit/',new_polling_unit,name = 'new_unit')

]


#!/usr/bin/env python3

import asyncio
import websockets
import ssl

async def hello(websocket, path):

    while True:
        m = await websocket.recv()
        print(m)

sc = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
# sc.check_hostname = False
# sc.verify_mode = ssl.CERT_NONE
sc.load_cert_chain('./resource/server.pem', './resource/server.key')

start_server = websockets.serve(hello, 'localhost', 8765, ssl=sc, subprotocols=[123, 32, "abc"])
# start_server = websockets.serve(hello, 'localhost', 8765)

asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()


import pytest
import sqlalchemy as sa
from fastapi.testclient import TestClient
from sqlalchemy.orm import Session, sessionmaker

from app import crud, models, schemas
from app.api.deps import get_db
from app.db.database import Base
from app.main import app

# TODO fix this to use database.py stuff
SQLALCHEMY_DATABASE_URL = "sqlite:///./test.db"

engine = sa.create_engine(
    SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}, future=True
)
TestingSessionLocal = sessionmaker(
    autocommit=False, autoflush=False, bind=engine, future=True
)

# Set up the database once
Base.metadata.drop_all(bind=engine)
Base.metadata.create_all(bind=engine)

_db = TestingSessionLocal()

class TestData:
    USER: models.User = crud.user.create(
        _db, obj_in=schemas.UserCreate(email="gert@bert.nl")
    )
    EXER_PARA_1: models.ExerciseParameter = crud.exercise_parameter.create(
        _db, obj_in=schemas.ExerciseParameterCreate(name="mm")
    )
    EXER_PARA_2: models.ExerciseParameter = crud.exercise_parameter.create(
        _db, obj_in=schemas.ExerciseParameterCreate(name="kg")
    )
    EXERCISE_1: models.Exercise = crud.exercise.save(
        _db,
        db_obj=models.Exercise(
            name="Moeilijke Oefening", parameters=[EXER_PARA_1, EXER_PARA_2]
        ),
    )
    EXERCISE_2: models.Exercise = crud.exercise.save(
        _db,
        db_obj=models.Exercise(name="Makkelijke Oefening", parameters=[EXER_PARA_1]),
    )

# These two event listeners are only needed for sqlite for proper
# SAVEPOINT / nested transaction support. Other databases like postgres
# don't need them.
# From: https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#serializable-isolation-savepoints-transactional-ddl # noqa
@sa.event.listens_for(engine, "connect")
def do_connect(dbapi_connection, connection_record):
    # disable pysqlite's emitting of the BEGIN statement entirely.
    # also stops it from emitting COMMIT before any DDL.
    dbapi_connection.isolation_level = None
    # dbapi_connection.execute("PRAGMA foreign_keys=ON;")

@sa.event.listens_for(engine, "begin")
def do_begin(conn):
    # emit our own BEGIN
    conn.exec_driver_sql("BEGIN")

# This fixture is the main difference to before. It creates a nested
# transaction, recreates it when the application code calls session.commit
# and rolls it back at the end.
# Based on: https://docs.sqlalchemy.org/en/14/orm/session_transaction.html#joining-a-session-into-an-external-transaction-such-as-for-test-suites # noqa
@pytest.fixture()
def db() -> Session:
    connection = engine.connect()
    transaction = connection.begin()
    session = TestingSessionLocal(bind=connection)

    # Begin a nested transaction (using SAVEPOINT).
    nested = connection.begin_nested()

    # If the application code calls session.commit, it will end the nested
    # transaction. Need to start a new one when that happens.
    @sa.event.listens_for(session, "after_transaction_end")
    def end_savepoint(session, transaction):
        nonlocal nested
        if not nested.is_active:
            nested = connection.begin_nested()

    yield session

    # Rollback the overall transaction, restoring the state before the test ran.
    session.close()
    transaction.rollback()
    connection.close()

# A fixture for the fastapi test client which depends on the
# previous session fixture. Instead of creating a new session in the
# dependency override as before, it uses the one provided by the
# session fixture.
@pytest.fixture()
def client(db) -> TestClient:
    def override_get_db():
        yield db

    app.dependency_overrides[get_db] = override_get_db
    yield TestClient(app)
    del app.dependency_overrides[get_db]