""" Tests knockpy.knockoffs and knockpy.smatrix modules"""
import warnings
import numpy as np
import scipy as sp
import unittest

# for regular pytest calls
try:
    from .context import knockpy
# for running directly with python
except ImportError:
    from context import knockpy
from knockpy import dgp, utilities, mac, mrc, smatrix, knockoffs

try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
try:
    import choldate
    CHOLDATE_AVAILABLE = True
except:
    CHOLDATE_AVAILABLE = False
try:
    import pydsdp
    DSDP_AVAILABLE = True
except:
    DSDP_AVAILABLE = False

class CheckSMatrix(unittest.TestCase):

    # Helper function
    def check_S_properties(self, V, S, groups):

        # Test PSD-ness of S
        min_S_eig = np.linalg.eigh(S)[0].min()
        self.assertTrue(
            min_S_eig > 0,
            f"S matrix is not positive semidefinite: mineig is {min_S_eig}",
        )

        # Test PSD-ness of 2V - S
        min_diff_eig = np.linalg.eigh(2 * V - S)[0].min()
        self.assertTrue(
            min_diff_eig > 0,
            f"2Sigma-S matrix is not positive semidefinite: mineig is {min_diff_eig}",
        )

        # Calculate conditional knockoff matrix
        invV = utilities.chol2inv(V)
        invV_S = np.dot(invV, S)
        Vk = 2 * S - np.dot(S, invV_S)

        # Test PSD-ness of the conditional knockoff matrix
        min_Vk_eig = np.linalg.eigh(Vk)[0].min()
        self.assertTrue(
            min_Vk_eig > 0,
            f"conditional knockoff matrix is not positive semidefinite: mineig is {min_Vk_eig}",
        )

        # Test that S is just a block matrix
        p = V.shape[0]
        S_test = np.zeros((p, p))
        for j in np.unique(groups):

            # Select subset of S
            inds = np.where(groups == j)[0]
            full_inds = np.ix_(inds, inds)
            group_S = S[full_inds]

            # Fill only in this subset of S
            S_test[full_inds] = group_S

        # return
        np.testing.assert_almost_equal(
            S_test,
            S,
            decimal=5,
            err_msg="S matrix is not a block matrix of the correct shape",
        )

class TestEquicorrelated(CheckSMatrix):
    """ Tests equicorrelated knockoffs and related functions """

    def test_eigenvalue_calculation(self):

        # Test to make sure non-group and group versions agree
        # (in the case of no grouping)
        p = 100
        groups = np.arange(0, p, 1) + 1
        for rho in [0, 0.3, 0.5, 0.7]:
            V = np.zeros((p, p)) + rho
            for i in range(p):
                V[i, i] = 1
            expected_gamma = min(1, 2 * (1 - rho))
            gamma = mac.calc_min_group_eigenvalue(Sigma=V, groups=groups,)
            np.testing.assert_almost_equal(
                gamma,
                expected_gamma,
                decimal=3,
                err_msg="calc_min_group_eigenvalue calculates wrong eigenvalue",
            )

        # Test non equicorrelated version
        V = np.random.randn(p, p)
        V = np.dot(V.T, V) + 0.1 * np.eye(p)
        V = utilities.cov2corr(V)
        expected_gamma = min(1, 2 * np.linalg.eigh(V)[0].min())
        gamma = mac.calc_min_group_eigenvalue(Sigma=V, groups=groups)
        np.testing.assert_almost_equal(
            gamma,
            expected_gamma,
            decimal=3,
            err_msg="calc_min_group_eigenvalue calculates wrong eigenvalue",
        )

    def test_equicorrelated_construction(self):

        # Test S matrix construction
        p = 100
        groups = np.arange(0, p, 1) + 1
        V = np.random.randn(p, p)
        V = np.dot(V.T, V) + 0.1 * np.eye(p)
        V = utilities.cov2corr(V)

        # Expected construction
        expected_gamma = min(1, 2 * np.linalg.eigh(V)[0].min())
        expected_S = expected_gamma * np.eye(p)

        # Equicorrelated
        S = mac.solve_equicorrelated(Sigma=V, groups=groups)

        # Test to make sure the answer is expected
        np.testing.assert_almost_equal(
            S,
            expected_S,
            decimal=3,
            err_msg="calc_min_group_eigenvalue calculates wrong eigenvalue",
        )

    def test_psd(self):

        # Test S matrix construction
        p = 100
        V = np.random.randn(p, p)
        V = np.dot(V.T, V) + 0.1 * np.eye(p)
        V = utilities.cov2corr(V)

        # Create random groups
        groups = np.random.randint(1, p, size=(p))
        groups = utilities.preprocess_groups(groups)
        S = mac.solve_equicorrelated(Sigma=V, groups=groups)

        # Check S properties
        self.check_S_properties(V, S, groups)

class TestSDP(CheckSMatrix):
    """ Tests an easy case of SDP and ASDP """

    def test_easy_sdp(self):

        # Test non-group SDP first
        n = 200
        p = 50
        X, _, _, _, corr_matrix, groups = dgp.block_equi_graph(
            n=n, p=p, gamma=0.3
        )

        # S matrix
        trivial_groups = np.arange(0, p, 1) + 1
        S_triv = smatrix.compute_smatrix(
            Sigma=corr_matrix, groups=trivial_groups, method="sdp", verbose=True,
        )
        np.testing.assert_array_almost_equal(
            S_triv,
            np.eye(p),
            decimal=2,
            err_msg="solve_group_SDP does not produce optimal S matrix (blockequi graphs)",
        )
        self.check_S_properties(corr_matrix, S_triv, trivial_groups)

        # Repeat for gaussian_knockoffs method
        ksampler = knockoffs.GaussianSampler(
            X=X, Sigma=corr_matrix, groups=trivial_groups, verbose=False, method="sdp",
        )
        S_triv2 = ksampler.fetch_S()

        np.testing.assert_array_almost_equal(
            S_triv2,
            np.eye(p),
            decimal=2,
            err_msg="solve_group_SDP does not produce optimal S matrix (blockequi graphs)",
        )
        self.check_S_properties(corr_matrix, S_triv2, trivial_groups)

        # Test slightly harder case
        _, _, _, _, expected_out, _ = dgp.block_equi_graph(n=n, p=p, gamma=0)
        ksampler = knockoffs.GaussianSampler(
            X=X, Sigma=corr_matrix, groups=groups, verbose=False, method="sdp"
        )
        S_harder = ksampler.fetch_S()
        np.testing.assert_almost_equal(
            S_harder,
            expected_out,
            decimal=2,
            err_msg="solve_group_SDP does not produce optimal S matrix (blockequi graphs)",
        )
        self.check_S_properties(corr_matrix, S_harder, groups)

        # Repeat for ASDP
        ksampler = knockoffs.GaussianSampler(
            X=X,
            Sigma=corr_matrix,
            groups=groups,
            method="ASDP",
            verbose=False,
            max_block=10,
        )
        S_harder_ASDP = ksampler.fetch_S()
        np.testing.assert_almost_equal(
            S_harder_ASDP,
            expected_out,
            decimal=2,
            err_msg="solve_group_ASDP does not produce optimal S matrix (blockequi graphs)",
        )
        self.check_S_properties(corr_matrix, S_harder_ASDP, groups)

    def test_equicorr_SDP(self):

        # Test non-group SDP on equicorrelated cov matrix
        p = 100
        rhos = [0.6, 0.8, 0.9]
        solvers = ['sdp', 'cd']
        scales = [1, 5]
        for scale in scales:
            for rho in rhos:
                V = scale * rho * np.ones((p, p)) + scale * (1 - rho) * np.eye(p)
                for solver in solvers:
                    S = smatrix.compute_smatrix(Sigma=V, method="sdp", solver=solver, verbose=True)
                    expected = (2 - 2 * rho) * np.eye(p)
                    np.testing.assert_almost_equal(
                        S / scale,
                        expected,
                        decimal=2,
                        err_msg=f"compute_smatrix produces incorrect S_SDP for equicorr, rho={rho}, scale={scale}, solver={solver}",
                    )

    def test_cd_SDP(self):
        """
        Test that CD / SDP solvers get similar answers in a hard problem
        """

        np.random.seed(110)
        V = utilities.cov2corr(dgp.ErdosRenyi(p=100, tol=1e-1))
        S_cd = smatrix.compute_smatrix(
            Sigma=V, method="sdp", solver="cd", verbose=True, mu=0.95, num_iter=100
        )
        mac_cd = np.diag(S_cd).mean()
        S_sdp = smatrix.compute_smatrix(Sigma=V, method="sdp", solver="sdp")
        mac_sdp = np.diag(S_sdp).mean()
        np.testing.assert_almost_equal(
            mac_sdp,
            mac_cd,
            decimal=2,
            err_msg=f"compute_smatrix mac differs between SDP/CD solvers ({mac_sdp}, {mac_cd}, respectively"
        )

    def test_sdp_tolerance(self):

        # Get graph
        np.random.seed(110)
        Q = dgp.ErdosRenyi(p=50, tol=1e-1)
        V = utilities.cov2corr(utilities.chol2inv(Q))
        groups = np.concatenate([np.zeros(10) + j for j in range(5)]) + 1
        groups = groups.astype("int32")

        # Solve SDP
        for tol in [1e-3, 0.01, 0.02]:
            S = smatrix.compute_smatrix(
                Sigma=V,
                groups=groups,
                method="sdp",
                objective="pnorm",
                num_iter=10,
                tol=tol,
            )
            G = np.hstack([np.vstack([V, V - S]), np.vstack([V - S, V])])
            mineig = np.linalg.eig(G)[0].min()
            self.assertTrue(
                tol - mineig > -1 * tol / 10,
                f"sdp solver fails to control minimum eigenvalues: tol is {tol}, val is {mineig}",
            )
            self.check_S_properties(V, S, groups)

    def test_corrmatrix_errors(self):
        """ Tests that SDP raises informative errors when sigma is not scaled properly"""

        # Get graph
        np.random.seed(110)
        Q = dgp.ErdosRenyi(p=50, tol=1e-1)
        V = utilities.chol2inv(Q)
        groups = np.concatenate([np.zeros(10) + j for j in range(5)]) + 1
        groups = groups.astype("int32")

        # Helper function
        def SDP_solver():
            return mac.solve_group_SDP(V, groups)

        # Make sure the value error increases
        self.assertRaisesRegex(
            ValueError, "Sigma is not a correlation matrix", SDP_solver
        )

class TestCholupdate(CheckSMatrix):
    """
    Checks that the cholupdate function works correctly.
    """

    def test_cholupdate(self):

        # Generate complex covariance matrix and vector
        p = 500
        dgprocess = knockpy.dgp.DGP()
        dgprocess.sample_data(p=p, method='qer', delta=0.8)
        V = dgprocess.Sigma + np.eye(p)
        L = np.linalg.cholesky(V)
        np.testing.assert_array_almost_equal(
            V,
            np.dot(L, L.T),
            decimal=6,
            err_msg="numpy cholesky decomposition failed"
        )
        x = np.random.uniform(size=(p))

        # Test dense update
        V = V + np.outer(x, x)
        knockpy.mrc.cholupdate(L.T, x, add=True)
        np.testing.assert_array_almost_equal(
            V,
            np.dot(L, L.T),
            decimal=6,
            err_msg="mrc.cholupdate with add=True failed"
        )

        # Test sparse downdate
        x = np.zeros(p)
        x[5] = 0.2
        V = V - np.outer(x, x)
        knockpy.mrc.cholupdate(L.T, x, add=False)
        np.testing.assert_array_almost_equal(
            V,
            np.dot(L, L.T),
            decimal=6,
            err_msg="mrc.cholupdate with add=False failed"
        )

class TestDependencyWarnings(CheckSMatrix):
    """ Tests DSDP_warning / choldate_warning """

    def test_warnings_raised(self):

        print(f"torch={TORCH_AVAILABLE}, dsdp={DSDP_AVAILABLE}, choldate={CHOLDATE_AVAILABLE}")

        # Sigma
        p = 50
        rho = 0.8
        Sigma = rho*np.ones((p,p)) + (1-rho)*np.eye(p)

        for method, dependency_flag, dependency_name, suppress_kwargs in zip(
            ['mvr', 'maxent', 'sdp'],
            [CHOLDATE_AVAILABLE, CHOLDATE_AVAILABLE, DSDP_AVAILABLE],
            ["choldate", "choldate", "dsdp"],
            [{"choldate_warning":False}, {"choldate_warning":False}, {"dsdp_warning":False}]
        ):
            # Check default setting
            with warnings.catch_warnings(record=True) as w:
                warnings.simplefilter('always')
                smatrix.compute_smatrix(Sigma, method=method)
                if not dependency_flag:
                    self.assertTrue(
                        dependency_name in str(w[-1].message)
                    )
                else:
                    self.assertTrue(len(w) == 0)
            # Check that it does not trigger with correct flag
            with warnings.catch_warnings(record=True) as w:
                warnings.simplefilter('always')
                smatrix.compute_smatrix(Sigma, method=method, **suppress_kwargs)
                self.assertTrue(len(w) == 0, "warning still triggers with choldate_warning=False")

class TestMRCSolvers(CheckSMatrix):
    """ Tests the various MRC solvers / classes"""

    def test_scale_sqrt_S(self):
        """ Tests the function which scales sqrt S for PSGD solver"""

        if not TORCH_AVAILABLE:
            return None
        from knockpy import kpytorch

        # Construct covariance matrix
        p = 50
        rho = 0.8
        Sigma = np.zeros((p, p)) + rho
        Sigma += (1 - rho) * np.eye(p)
        groups = np.arange(1, p + 1, 1)
        init_blocks = [np.eye(1) for _ in range(p)]

        # Create model - this automatically scales the
        # initial blocks properly
        fk_precision_calc = kpytorch.mrcgrad.MVRLoss(Sigma, groups, init_S=init_blocks)
        # Check for proper scaling
        S = fk_precision_calc.pull_S().detach().numpy()
        expected = min(1, 2 - 2 * rho) * np.eye(p)
        np.testing.assert_almost_equal(
            S,
            expected,
            decimal=1,
            err_msg=f"Initial scaling fails, expected {expected} but got {S} for equicorrelated rho={rho}",
        )

    def test_group_sorting_error(self):
        """ Tests PSGD class raises error if the cov
        matrix/groups are improperly sorted"""

        if not TORCH_AVAILABLE:
            return None
        from knockpy import kpytorch

        # Groups and sigma
        p = 50
        Sigma = np.eye(p)
        groups = knockpy.utilities.preprocess_groups(np.random.randint(1, p + 1, p))

        # Try to initialize
        def init_unsorted_model():
            model = kpytorch.mrcgrad.MVRLoss(Sigma, groups)

        self.assertRaisesRegex(
            ValueError,
            "Sigma and groups must be sorted prior to input",
            init_unsorted_model,
        )

    def test_smoothing(self):
        """
        Smoothing is not required for this, but this is a nice check anyway.
        """
        p = 50
        smoothing = 0.1
        dgprocess = dgp.DGP()
        _, _, _, _, V = dgprocess.sample_data(method="partialcorr", rho=0.1,)
        S_MVR = mrc.solve_mvr(Sigma=V, smoothing=smoothing)
        # Not implemented yet
        S_SDP = mac.solve_SDP(Sigma=V, tol=1e-5)
        mvr_mean = np.diag(S_MVR).mean()
        sdp_mean = np.diag(S_SDP).mean()
        self.assertTrue(
            sdp_mean - mvr_mean < 1e-3,
            f"Highly smoothed S_MVR ({S_MVR}) too far from S_SDP ({S_SDP}) for equi partial corr",
        )

    def test_equicorrelated_soln(self):
        """ Tests that solvers yield expected
        solution for equicorrelated matrices """

        if not TORCH_AVAILABLE:
            return None
        from knockpy import kpytorch

        # Main constants
        p = 50
        groups = np.arange(1, p + 1, 1)
        smoothings = [0, 0.01]

        # Construct equicorrelated matrices
        rhos = [0.1, 0.3, 0.5, 0.7, 0.9]
        for rho in rhos:
            for smoothing in smoothings:
                for method in ["mvr", "mmi", "maxent"]:
                    # Construct Sigma
                    Sigma = np.zeros((p, p)) + rho
                    Sigma = Sigma + (1 - rho) * np.eye(p)

                    # Expected solution
                    opt_prop_rec = min(rho, 0.5)
                    max_S_val = min(1, 2 - 2 * rho)
                    expected = (1 - rho) * np.eye(p)

                    # Test optimizer
                    opt_S = kpytorch.mrcgrad.solve_mrc_psgd(
                        Sigma=Sigma,
                        groups=groups,
                        init_S=None,
                        smoothing=smoothing,
                        verbose=True,
                        tol=1e-5,
                        max_epochs=100,
                        line_search_iter=10,
                        lr=1e-2,
                    )
                    self.check_S_properties(Sigma, opt_S, groups)
                    np.testing.assert_almost_equal(
                        opt_S,
                        expected,
                        decimal=1,
                        err_msg=f"For equicorrelated cov rho={rho}, PSGD solver yields unexpected solution",
                    )

                # Test MVR coordinate descent optimizer
                opt_S = mrc.solve_mvr(Sigma=Sigma, smoothing=smoothing, verbose=True)
                self.check_S_properties(Sigma, opt_S, groups)
                np.testing.assert_almost_equal(
                    opt_S,
                    expected,
                    decimal=2,
                    err_msg=f"For equicorrelated cov rho={rho}, mvr_solver yields unexpected solution",
                )

                # Test maximum entropy coordinate descent optimizer
                if smoothing == 0:
                    opt_S = mrc.solve_mmi(
                        Sigma=Sigma, verbose=True
                    )
                    self.check_S_properties(Sigma, opt_S, groups)
                    np.testing.assert_almost_equal(
                        opt_S,
                        expected,
                        decimal=2,
                        err_msg=f"For equicorrelated cov rho={rho}, mmi_solver yields unexpected solution",
                    )

    def test_maxent(self):
        """ Both maxent/mmi work properly """
        # Sample data
        dgprocess = dgp.DGP()
        dgprocess.sample_data(p=50, method='ar1', a=3)

        # Check solve_maxent/solve_mmi
        np.random.seed(110)
        S_ME = smatrix.compute_smatrix(dgprocess.Sigma, method='maxent')
        np.random.seed(110)
        S_MMI = smatrix.compute_smatrix(dgprocess.Sigma, method='mmi')
        np.testing.assert_array_almost_equal(
            S_ME, S_MMI, decimal=3, err_msg=f"compute_smatrix yields diff answers for mmi, maxent"
        )

        # Check solve_maxent/solve_mmi
        np.random.seed(110)
        S_ME = mrc.solve_maxent(dgprocess.Sigma)
        np.random.seed(110)
        S_MMI = mrc.solve_mmi(dgprocess.Sigma)
        np.testing.assert_array_almost_equal(
            S_ME, S_MMI, decimal=3, err_msg=f"solve_maxent and solve_mmi yield different answers"
        )

        # Check maxent_loss/mmi_loss
        L_ME = mrc.maxent_loss(dgprocess.Sigma, S_ME)
        L_MMI = mrc.mmi_loss(dgprocess.Sigma, S_MMI)
        np.testing.assert_almost_equal(
            L_ME, L_MMI, decimal=3, err_msg=f"maxent_loss and mmi_loss yield different answers"
        )

    def test_equicorrelated_soln_recycled(self):
        """ Correct solutions on equicorrelated matrix with recycling """

        # Potentially also test the PSGD solver if torch is available
        if TORCH_AVAILABLE:
            from knockpy import kpytorch

        # Main constants
        p = 50
        groups = np.arange(1, p + 1, 1)

        # Test separately with the recycling proportion param
        rhos = [0.1, 0.3, 0.5, 0.7, 0.9]
        true_rec_props = [0.5, 0.25, 0.8, 0.5, 0.5]
        for true_rec_prop, rho in zip(true_rec_props, rhos):

            # Construct Sigma
            Sigma = np.zeros((p, p)) + rho
            Sigma = Sigma + (1 - rho) * np.eye(p)

            # Expected solution
            opt_prop_rec = min(rho, 0.5)
            max_S_val = min(1, 2 - 2 * rho)
            normal_opt = (1 - opt_prop_rec) * max_S_val
            new_opt = min(2 - 2 * rho, normal_opt / (1 - true_rec_prop))
            expected = new_opt * np.eye(p)

            # Test PSGD optimizer
            if TORCH_AVAILABLE:
                opt_S = kpytorch.mrcgrad.solve_mrc_psgd(
                    Sigma=Sigma,
                    groups=groups,
                    init_S=None,
                    rec_prop=true_rec_prop,
                    tol=1e-5,
                    max_epochs=100,
                    line_search_iter=10,
                )
                self.check_S_properties(Sigma, opt_S, groups)
                np.testing.assert_almost_equal(
                    opt_S,
                    expected,
                    decimal=2,
                    err_msg=f"For equicorrelated cov rho={rho} rec_prop={true_rec_prop}, MVR PSDSolver returns {opt_S}, expected {expected}",
                )

            # Coord descent solver
            np.random.seed(110)
            opt_S = mrc.solve_mvr(
                Sigma=Sigma, rej_rate=true_rec_prop, tol=1e-5, verbose=True,
            )
            self.check_S_properties(Sigma, opt_S, groups)
            np.testing.assert_almost_equal(
                opt_S,
                expected,
                decimal=2,
                err_msg=f"For equicorrelated cov rho={rho} rec_prop={true_rec_prop}, MVR CD solver returns {opt_S}, expected {expected}",
            )

    def test_complex_solns(self):
        """
        Check the solution of the various solvers
        for non-grouped knockoffs.
        """

        # Check availability
        if not TORCH_AVAILABLE:
            return None
        from knockpy import kpytorch

        np.random.seed(110)
        p = 100
        methods = ["ar1", "ver"]
        groups = np.arange(1, p + 1, 1)
        for method in methods:
            dgprocess = dgp.DGP()
            _, _, _, _, Sigma = dgprocess.sample_data(method=method, p=p)

            # Use SDP as baseline
            init_S = knockpy.mac.solve_group_SDP(Sigma, groups)
            sdp_mvr_loss = mrc.mvr_loss(Sigma, init_S)

            # Apply gradient solver
            opt_S = kpytorch.mrcgrad.solve_mrc_psgd(
                Sigma=Sigma,
                groups=groups,
                init_S=init_S,
                tol=1e-5,
                max_epochs=100,
                line_search_iter=10,
            )
            psgd_mvr_loss = mrc.mvr_loss(Sigma, opt_S)

            # Check S matrix
            self.check_S_properties(Sigma, opt_S, groups)
            # Check new loss < init_loss
            self.assertTrue(
                psgd_mvr_loss <= sdp_mvr_loss,
                msg=f"For {method}, PSGD solver has higher loss {psgd_mvr_loss} v. sdp {sdp_mvr_loss}",
            )

            # MVR solver outperforms PSGD
            opt_S_mvr = mrc.solve_mvr(Sigma=Sigma)
            self.check_S_properties(Sigma, opt_S_mvr, groups)
            cd_mvr_loss = mrc.mvr_loss(Sigma, opt_S_mvr)
            self.assertTrue(
                cd_mvr_loss <= psgd_mvr_loss,
                msg=f"For {method}, coord descent MVR solver has higher loss {cd_mvr_loss} v. PSGD {psgd_mvr_loss}",
            )

            # mmi solver outperforms PSGD
            print(Sigma)
            print(np.linalg.eigh(Sigma)[0].min())
            opt_S_mmi = mrc.solve_mmi(Sigma=Sigma)
            print(opt_S_mmi)
            print(opt_S)
            self.check_S_properties(Sigma, opt_S_mmi, groups)
            cd_mmi_loss = mrc.mmi_loss(Sigma, opt_S_mmi)
            psgd_mmi_loss = mrc.mmi_loss(Sigma, opt_S)
            self.assertTrue(
                cd_mmi_loss <= psgd_mmi_loss,
                msg=f"For {method}, coord descent mmi solver has higher loss {cd_mmi_loss} v. PSGD {psgd_mmi_loss}",
            )

    def test_cd_group_mvr(self):

        np.random.seed(110)
        p = 100
        for method in ['ar1', 'ver', 'qer']:
            # Sample data
            dgprocess = knockpy.dgp.DGP()
            dgprocess.sample_data(method=method, p=p)
            Sigma = dgprocess.Sigma

            # Sanity check: grouped solver gets right soln for ungrouped
            S1 = mrc._solve_mvr_grouped(
                Sigma, groups=np.arange(1, p+1), converge_tol=1e-2, verbose=True
            )
            S2 = mrc.solve_mvr(Sigma)
            Sdiff = S1 - S2
            mean_diff = np.abs(Sdiff[Sdiff != 0]).mean()
            err_msg = f"For no groups, method={method}, avg diff. between grouped/ungrouped solvers is {mean_diff} > 1e-2"
            err_msg = f"grouped solver: {S1} \n ungrouped solver: {S2} \n" + err_msg
            self.assertTrue(
                mean_diff < 1e-2, err_msg
            )

            # Check that adding groups lowers loss
            loss_ungrouped = mrc.mvr_loss(Sigma, S2)
            groups = np.around(np.arange(1, p+1)/2)
            Sgroup = mrc.solve_mvr(Sigma, groups)
            loss_grouped = mrc.mvr_loss(Sigma, Sgroup)
            self.assertTrue(
                loss_grouped < loss_ungrouped,
                f"Grouped loss ({loss_grouped}) > ungrouped loss ({loss_ungrouped}) for method={method}"
            )

    def test_complex_group_solns(self):
        """
        Check the solutions of the PSGD solver
        for group knockoffs.
        """

        if not TORCH_AVAILABLE:
            return None
        from knockpy import kpytorch

        # Construct graph + groups
        np.random.seed(110)
        p = 50
        groups = knockpy.utilities.preprocess_groups(np.random.randint(1, p + 1, p))
        for method in ["ar1", "ver"]:
            dgprocess = dgp.DGP()
            _, _, _, _, Sigma = dgprocess.sample_data(method=method, p=p,)

            # Use SDP as baseline
            init_S = knockpy.mac.solve_group_SDP(Sigma, groups)
            init_loss = mrc.mvr_loss(Sigma, init_S)

            # Apply gradient solver
            opt_S = kpytorch.mrcgrad.solve_mrc_psgd(
                Sigma=Sigma,
                groups=groups,
                init_S=init_S,
                tol=1e-5,
                max_epochs=100,
                line_search_iter=10,
            )
            psgd_loss = mrc.mvr_loss(Sigma, opt_S)

            # Check S matrix
            self.check_S_properties(Sigma, opt_S, groups)
            # Check new loss < init_loss
            self.assertTrue(
                psgd_loss <= init_loss,
                msg=f"For {method}, PSGD solver has higher loss {psgd_loss} v. sdp {init_loss}",
            )

    def test_equi_ciknock_solution(self):
        """
        Check ciknockoff solution
        """
        p = 500
        rho = 0.6
        # 1. Block equicorrelated
        dgprocess = dgp.DGP()
        _, _, _, _, Vblock = dgprocess.sample_data(
            p=p, method="blockequi", gamma=0, rho=rho, block_size=2
        )
        S_CI = mrc.solve_ciknock(Vblock)
        np.testing.assert_almost_equal(
            S_CI,
            (1 - rho ** 2) * np.eye(p),
            2,
            "S_CI is incorrect for block-equicorrelated with blocksize 2",
        )
        # 2. Equicorelated
        dgprocess = dgp.DGP()
        _, _, _, _, V = dgprocess.sample_data(
            p=p, method="blockequi", gamma=1, rho=rho
        )
        S_CI = mrc.solve_ciknock(V)
        np.testing.assert_almost_equal(
            S_CI, (1 - rho) * np.eye(p), 2, "S_CI is incorrect for equicorrelated"
        )

class TestBlockdiagApprx(CheckSMatrix):
    def test_divide_comp(self):
        """ Check that when we divide computation, our
        final blocks do not exceed the specified block size."""

        p = 120
        for method in ["blockequi", "ar1", "ver", "qer"]:
            dgprocess = dgp.DGP()
            _, _, _, _, V = dgprocess.sample_data(p=p, method=method)
            for max_block in [3, 12, 15, 54, 76]:
                blocks = smatrix.divide_computation(V, max_block)
                block_sizes = utilities.calc_group_sizes(blocks)
                self.assertTrue(
                    block_sizes.max() <= max_block,
                    msg=f"Returned blocks have max size {block_sizes.max()} > max_block {max_block}",
                )

    def test_blockdiag_approx_on_blockdiag(self):

        # Check that blockdiag approx yields same solution
        # for block-diagonal matrix.
        p = 250
        rho = 0.3
        gamma = 0
        groups = np.arange(1, p + 1, 1)
        dgprocess = dgp.DGP()
        _, _, _, _, V = dgprocess.sample_data(
            p=p, method="blockequi", gamma=gamma, rho=rho
        )

        for method in ["mvr", "maxent", "sdp"]:
            # Check that S_approx is valid
            S_approx = smatrix.compute_smatrix(V, method=method, max_block=10)
            self.check_S_properties(V, S_approx, groups)

            # In this case, should be the same as S_exact
            S_exact = smatrix.compute_smatrix(V, method=method)
            np.testing.assert_array_almost_equal(
                S_approx,
                S_exact,
                2,
                err_msg="Blockdiag approximation yields incorrect answer for blockdiag Sigma",
            )

    def test_blockdiag_approx_on_ar1(self):

        # Check that blockdiag approx yields a good soln on AR1
        p = 150
        a = 1
        b = 1
        # Create trivial and nontrivial groups
        triv_groups = np.arange(1, p + 1, 1)
        nontriv_groups = np.around(triv_groups / 2 + 0.01)
        nontriv_groups = utilities.preprocess_groups(nontriv_groups)
        # Sample data
        np.random.seed(110)
        dgprocess = dgp.DGP()
        _, _, _, _, V = dgprocess.sample_data(p=p, method="ar1", a=a, b=b)
        for groups in [triv_groups, nontriv_groups]:
            for method in ["mvr", "maxent", "mmi", "sdp"]:

                # Note we can't fit group mvr/maxent without pytorch
                nontriv_group_flag = np.all(groups == nontriv_groups)
                if not TORCH_AVAILABLE and nontriv_group_flag:
                    continue

                # Check that S_approx is valid
                S_approx = smatrix.compute_smatrix(
                    V, method=method, groups=groups, max_block=50
                )
                self.check_S_properties(V, S_approx, groups)

                # For exponentially decaying offdiags, should be quite similar
                # This seems to not be true for groups, so skip that for now...
                # For now, skip grouped mvr/mmi
                if nontriv_group_flag:
                    continue
                S_exact = smatrix.compute_smatrix(V, method=method, groups=groups)
                diff = np.abs(S_approx - S_exact)
                mean_diff = diff[S_exact != 0].mean()
                expected = 1e-2 if CHOLDATE_AVAILABLE else 5e-2
                identifier = f"for method={method}, groups={groups}"
                self.assertTrue(
                    mean_diff < expected,
                    msg=f"Blockdiag apprx is {mean_diff} > {expected} on avg. away from exact soln {identifier}",
                )

class TestFactorApprx(CheckSMatrix):

    def test_agreement_on_factor_models(self):

        # Random factor models
        np.random.seed(110)
        ks = [1, 10, 20, 50]
        p = 150
        for k in ks:
            for method in ['mvr', 'maxent', 'sdp']:
                D = np.random.uniform(low=0.01, high=1, size=(p,))
                U = np.random.randn(p, k) / np.sqrt(k)
                # Rescale to correlation matrix
                diag_Sigma = D + (np.power(U, 2)).sum(axis=1)
                D = D / diag_Sigma
                U = U / np.sqrt(diag_Sigma).reshape(-1, 1)
                Sigma = np.diag(D) + np.dot(U, U.T)
                # Check the factor method agrees with reg. method
                S1 = smatrix.compute_smatrix(
                    Sigma=None, D=D, U=U, method=method, how_approx='factor'
                )
                S2 = smatrix.compute_smatrix(
                   Sigma=Sigma, method=method, solver='cd', tol=0
                )
                np.testing.assert_array_almost_equal(
                    S1, S2, decimal=2,
                    err_msg=f"factored/non-factored versions for method={method}, k={k} do not agree"
                )

class CheckValidKnockoffs(unittest.TestCase):
    def check_valid_mxknockoffs(self, X, mu=None, Sigma=None, msg="", **kwargs):

        # S matrix
        ksampler = knockoffs.GaussianSampler(
            X=X, mu=mu, Sigma=Sigma, verbose=True, **kwargs
        )
        S = ksampler.fetch_S()
        Xk = ksampler.sample_knockoffs()

        # Test knockoff mean
        if mu is None:
            mu = X.mean(axis=0)
        outmsg = "Knockoffs have incorrect means "
        outmsg += f"for MX knockoffs for {msg}"
        np.testing.assert_array_almost_equal(Xk.mean(axis=0), mu, 2, outmsg)

        # Sigma should be
        if Sigma is None:
            Sigma, _ = utilities.estimate_covariance(X, tol=1e-2)

        # Also rescale X/Xk so S makes sense
        scale = np.sqrt(np.diag(Sigma))
        X = X / scale.reshape(1, -1)
        if mu is None:
            mu = X.mean(axis=0)
        else:
            mu = mu / scale
        Xk = Xk / scale.reshape(1, -1)
        Sigma = Sigma / np.outer(scale, scale)
        S = S / np.outer(scale, scale)

        # Empirical FK correlation matrix
        features = np.concatenate([X, Xk], axis=1)
        G_hat = np.cov(features.T)

        # Calculate population version
        G = np.concatenate(
            [np.concatenate([Sigma, Sigma - S]), np.concatenate([Sigma - S, Sigma])],
            axis=1,
        )

        # Test G has correct structure
        outmsg = f"Feature-knockoff cov matrix has incorrect values "
        outmsg += f"for MX knockoffs for {msg} graph "
        np.testing.assert_array_almost_equal(G_hat, G, 2, outmsg)

class TestKnockoffGen(CheckValidKnockoffs):
    """ Tests whether knockoffs have correct distribution empirically"""

    def test_method_parser(self):

        # Easiest test
        method1 = "hello"
        out1 = smatrix.parse_method(method1, None, None)
        self.assertTrue(
            out1 == method1, "parse method fails to return non-None methods"
        )

        # Default is mvr
        p = 1000
        groups = np.arange(1, p + 1, 1)
        out2 = smatrix.parse_method(None, groups, p)
        self.assertTrue(out2 == "mvr", "parse method fails to return mvr by default")

        # Otherwise SDP
        groups[-1] = 1
        out2 = smatrix.parse_method(None, groups, p)
        self.assertTrue(
            out2 == "sdp", "parse method fails to return SDP for grouped knockoffs"
        )

    def test_error_raising(self):

        # Generate data
        n = 100
        p = 100
        X, _, _, _, corr_matrix, groups = dgp.block_equi_graph(
            n=n, p=p, gamma=1, rho=0.8
        )
        S_bad = np.eye(p)

        def fdr_vio_knockoffs():
            ksampler = knockoffs.GaussianSampler(
                X=X, Sigma=corr_matrix, S=S_bad, verbose=False
            )
            ksampler.sample_knockoffs(check_psd=True)

        self.assertRaisesRegex(
            np.linalg.LinAlgError,
            "meaning FDR control violations are extremely likely",
            fdr_vio_knockoffs,
        )

        # Test FX knockoff violations
        def fx_knockoffs_low_n():
            knockoffs.FXSampler(
                X=X, Sigma=corr_matrix, S=None,
            )

        self.assertRaisesRegex(
            np.linalg.LinAlgError,
            "FX knockoffs can't be generated with n",
            fx_knockoffs_low_n,
        )

    def test_consistency_of_inferring_sigma(self):
        """ Checks that the same knockoffs are produced
        whether you infer the covariance matrix first and
        pass it to the gaussian_knockoffs generator, or
        you let the generator do the work for you
        """

        n = 25
        p = 300
        rho = 0.5
        dgprocess = dgp.DGP()
        X, _, _, _, _ = dgprocess.sample_data(n=n, p=p, rho=rho, method="AR1")

        # Method 1: infer cov first
        V, _ = utilities.estimate_covariance(X, tol=1e-2)
        np.random.seed(110)
        Xk1 = knockoffs.GaussianSampler(
            X=X, Sigma=V, method="sdp"
        ).sample_knockoffs()

        # Method 2: Infer during
        np.random.seed(110)
        Xk2 = knockoffs.GaussianSampler(
            X=X, method="sdp"
        ).sample_knockoffs()
        np.testing.assert_array_almost_equal(
            Xk1, Xk2, 5, err_msg="Knockoff gen is inconsistent"
        )

    def test_MX_knockoff_dist(self):

        # Test knockoff construction for mvr and SDP
        # on equicorrelated matrices
        np.random.seed(110)
        n = 100000
        copies = 3
        p = 5

        # Check with a non-correlation matrix
        V = 4 * dgp.AR1(p=p, rho=0.5)
        mu = np.random.randn(p)
        print(f"true mu: {mu}")
        dgprocess = dgp.DGP(mu=mu, Sigma=V)
        X, _, _, _, _ = dgprocess.sample_data(n=n, p=p,)
        print(f"X mean: {X.mean(axis=0)}")

        # Check validity for oracle cov matrix
        self.check_valid_mxknockoffs(X, mu=mu, Sigma=V, msg=f"ORACLE 4*AR1(rho=0.5)")

        # Check validity for estimated cov matrix
        print("Here now")
        self.check_valid_mxknockoffs(X, msg=f"ESTIMATED 4*AR1(rho=0.5)")

        # Check for many types of data
        for rho in [0.1, 0.9]:
            for gamma in [0.5, 1]:
                for method in ["mvr", "sdp"]:

                    mu = 10 * np.random.randn(p)
                    X, _, _, _, corr_matrix, _ = dgp.block_equi_graph(
                        n=n, p=p, gamma=gamma, rho=rho, mu=mu
                    )

                    # Check validity for oracle correlation matrix
                    self.check_valid_mxknockoffs(
                        X,
                        mu=mu,
                        Sigma=corr_matrix,
                        msg=f"blockequi graph, rho = {rho}, gamma = {gamma}",
                    )

                    # Check validity for estimation
                    self.check_valid_mxknockoffs(
                        X,
                        msg=f"ESTIMATED blockequi graph, rho = {rho}, gamma = {gamma}",
                    )

    def test_FX_knockoff_dist(self):
        # Test knockoff construction for mvr and SDP
        # on equicorrelated matrices
        n = 500
        p = 5
        for rho in [0.1, 0.9]:
            for gamma in [0.5, 1]:
                for method in ["mvr", "sdp"]:
                    # X values
                    X, _, _, _, corr_matrix, _ = dgp.block_equi_graph(
                        n=n, p=p, gamma=gamma, rho=rho
                    )
                    # S matrix
                    trivial_groups = np.arange(0, p, 1) + 1
                    ksampler = knockoffs.FXSampler(X=X, method=method, verbose=False)
                    Xk = ksampler.sample_knockoffs()
                    S = ksampler.fetch_S()

                    # Compute empirical (scaled) cov matrix
                    features = np.concatenate([X, Xk], axis=1)
                    G_hat = np.dot(features.T, features)

                    # Calculate what this should be
                    Sigma = np.dot(X.T, X)
                    G = np.concatenate(
                        [
                            np.concatenate([Sigma, Sigma - S]),
                            np.concatenate([Sigma - S, Sigma]),
                        ],
                        axis=1,
                    )

                    # Test G has correct structure
                    msg = f"Feature-knockoff cov matrix has incorrect values "
                    msg += f"for blockequi graph, FX knockoffs, rho = {rho}, gamma = {gamma}"
                    np.testing.assert_array_almost_equal(G_hat, G, 5, msg)

    def test_scaling_consistency(self):
        """
        Checks whether if we first calculate S
        and then pass it back into the knockoff
        generator, we'll get the same answer back.
        """

        p = 100
        n = 300
        dgprocess = dgp.DGP()
        X, y, beta, Q, V = dgprocess.sample_data(
            method="qer", p=p, n=n, coeff_size=0.5, sparsity=0.5,
        )
        ksampler1 = knockoffs.FXSampler(X=X,)
        S1 = ksampler1.fetch_S()
        ksampler2 = knockoffs.FXSampler(X=X, S=S1,)
        S2 = ksampler2.fetch_S()
        np.testing.assert_array_almost_equal(
            S1,
            S2,
            decimal=6,
            err_msg=f"Repeatedly passing S into/out of knockoff gen yields inconsistencies",
        )

if __name__ == "__main__":
    import pytest
    import sys
    pytest.main(sys.argv)
    #unittest.main()


import os

def cls():
    os.system('cls' if os.name=='nt' else 'clear')

player_dict = {"p1":"","p2":""}
field_list = [" "," "," "," "," "," "," "," "," "," "]
marks_checklist = [field_list[0:9:4],field_list[2:7:2],field_list[0:7:3],field_list[1:8:3],
                   field_list[2:9:3],field_list[0:3],field_list[3:6],field_list[6:9]]
active_player = ""
X_points = 0
O_points = 0
moves = 9
gameon_choice = "wrong"
insert_coin = "error"

def restart_game():
    global gameon_choice
    global field_list
    global player_dict
    global active_player
    global moves

    active_player = ""
    field_list = [" "," "," "," "," "," "," "," "," "," "]
    moves = 9
    gameon_choice = "wrong"

def end_game():
    global gameon_choice
    global field_list
    global player_dict
    global active_player
    global X_points
    global O_points
    global moves
    global insert_coin
    player_dict = {"p1":"","p2":""}
    active_player = ""
    field_list = [" "," "," "," "," "," "," "," "," "," "]
    X_points = 0
    O_points = 0
    moves = 9
    gameon_choice = "wrong"
    insert_coin = "error"

marks_checklist = [field_list[1:10:4],field_list[3:8:2],field_list[1:8:3],field_list[2:9:3],
                   field_list[3:10:3],field_list[1:4],field_list[4:7],field_list[7:10]]

def win_check():
    global X_points
    global O_points
    marks_checklist = [field_list[1:10:4],field_list[3:8:2],field_list[1:8:3],field_list[2:9:3],
                   field_list[3:10:3],field_list[1:4],field_list[4:7],field_list[7:10]]
    for group in marks_checklist:

        if group == ['X','X','X']:
            X_points +=1

            print("X wins")
            return True
        elif group == ['O', 'O', 'O']:
            O_points +=1

            print( "O wins")
            return True
    else:
        return False

def tie_check():
    if moves == 0:
        print("TIE")
        return True
    else:
        return False

def display_board():

    empty_line = "    |  "" |"
    hr_line = "-------------"
    num_line79 = "  ""7 ""|"" 8 ""|"" 9"" "
    num_line46 = "  ""4 ""|"" 5 ""|"" 6"" "
    num_line13 = "  ""1 ""|"" 2 ""|"" 3"" "

    empty_line_aligned = empty_line.rjust(30)
    hr_line_aligned = hr_line.rjust(30)
    num_line79_aligned = num_line79.rjust(29)
    num_line46_aligned = num_line46.rjust(29)
    num_line13_aligned = num_line13.rjust(29)

    print("==============================================================")
    print("                                                             ")
    print("                                                             ")
    print("                                                             ")
    print("    |  "" |",empty_line_aligned)
    print(" ",field_list[7] ,"|",field_list[8] ,"|",field_list[9] ," ",num_line79_aligned)
    print("    |  "" |",empty_line_aligned)
    print("-------------",hr_line_aligned)
    print("    |  "" |",empty_line_aligned)
    print(" ",field_list[4] ,"|",field_list[5] ,"|",field_list[6] ," ",num_line46_aligned)
    print("    |  "" |",empty_line_aligned)
    print("-------------",hr_line_aligned)
    print("    |  "" |",empty_line_aligned)
    print(" ",field_list[1] ,"|",field_list[2] ,"|",field_list[3] ," ",num_line13_aligned)
    print("    |  "" |",empty_line_aligned)
    print("                                                             ")
    print("                                                             ")
    print("                                                             ")
    print("==============================================================")

from IPython.display import clear_output

def choose_side():
    side_choice = "wrong"
    while side_choice != 'X' and side_choice != 'O':

        side_choice = input('Do you want to play X or O?: ')

        if side_choice != 'X' or 'x' and side_choice != 'O' or 'o':
            clear_output()
            print('Please choose X or O')

    return side_choice

def player_assignment():
    player_one = choose_side()
    player_two = ""
    if player_one == "X":
        player_two = "O"
    elif player_one == "O":
        player_two = "X"
    player_dict['p1'] = player_one
    player_dict['p2'] = player_two
    print("Player 1: ",player_one,"Player 2: ",player_two)
    return player_dict

import random
import time

def who_starts(player_dict):
    global active_player

    num = random.randint(1,100)
    print("")
    print("Starting draw who starts...")
    print("")
    time.sleep(1.5)
    if num % 2 == 0:
        print("Player 1 starts")
        active_player = player_dict['p1']
        return active_player
    else:
        print("Player 2 starts")
        active_player = player_dict['p2']
        return active_player

def whos_next():
    global active_player
    if active_player == 'O':
        active_player = 'X'
    elif active_player == 'X':
        active_player = 'O'

    return active_player

def player_field_choice():

    global active_player
    global moves
    occupied = False
    num_choice = "wrong"
    while num_choice.isdigit() == False or num_choice not in ['1','2','3','4','5','6','7','8','9'] or occupied == False:
        display_board()
        print("")
        print("X: ",X_points," O: ",O_points)
        print("Player ",active_player," turn")
        num_choice = input("Choose a number from 1 to 9:")

        if num_choice.isdigit() == False:
            clear_output()
            print("This is not an integer, please type a number from 1 to 9")

        elif num_choice not in ['1','2','3','4','5','6','7','8','9']:
                clear_output()
                print("Number out of range, please type a number from 1 to 9")

        elif field_list[int(num_choice)] == 'O' or field_list[int(num_choice)] == 'X':
            clear_output()
            print("This field is already taken, please choose other field.")
            occupied = False

        elif field_list[int(num_choice)] == ' ':
            occupied = True

    if active_player == 'X':
        clear_output()

        moves -= 1
        field_list[int(num_choice)] = "X"
        cls()
        display_board()

    elif active_player == 'O':
        clear_output()

        moves -= 1
        field_list[int(num_choice)] = "O"
        cls()
        display_board()

    active_player = whos_next()

def play_again():
    global gameon_choice
    while gameon_choice != 'y'and gameon_choice != 'n':
        gameon_choice = input('Do you want to play again? y/n')
        if gameon_choice != 'y'and gameon_choice != 'n':
            clear_output()
            print('Please type "y" for Yes or "n" for No.')

        return gameon_choice

def launch():
    global insert_coin
    while insert_coin != 'y' and insert_coin != 'n':
        insert_coin = input('Game start? y/n')
        if insert_coin != 'y' and insert_coin != 'n':
            clear_output()
            print('Please type "y" for Yes or "n" for No.')

def game_start():
    player_assignment()
    who_starts(player_dict)
    while win_check() == False and tie_check() == False:

        clear_output()
        player_field_choice()
    play_again()

def game_on():
    who_starts(player_dict)
    while win_check() == False and tie_check() == False:
        clear_output()
        player_field_choice()

    play_again()

def game():
    global X_points
    global O_points
    game_start()
    while gameon_choice == 'y':
        restart_game()
        game_on()
    if gameon_choice == 'n':
        if X_points > O_points:
            print("X: ",X_points," O: ",O_points)
            print("And the winner is Player X with ",X_points," points.")
            end_game()
        elif X_points < O_points:
            print("X: ",X_points," O: ",O_points)
            print("And the winner is Player O with ",O_points," points.")
            end_game()
        elif X_points == O_points:
            print("X: ",X_points," O: ",O_points)
            print("Draw, X got ",X_points," O got ",O_points)
            end_game()

def tic_tac_toe():
    print('Welcome to Tic Tac Toe!')
    launch()
    if insert_coin == 'y':
        game()
    elif insert_coin == 'n':
        print("GAME OVER")
        end_game()

tic_tac_toe()


# standard library
import os

# dash libs
import dash
from dash.dependencies import Input, Output, State, Event
import dash_core_components as dcc
import dash_html_components as html
import plotly.figure_factory as ff
import plotly.graph_objs as go

# pydata stack
import pandas as pd
from sqlalchemy import create_engine

from datetime import datetime as dt

import io
import flask
from flask import send_file, url_for, request
import urllib.parse as p
import xlsxwriter
import zipfile
import io
import pathlib
import numpy as np

# set params
#conn = create_engine(os.environ['DB_URI'])
conn = create_engine("sqlite:///atcTyres.db")
globalStockReport = pd.DataFrame(columns = ['plantCode', 'rmCode', 'calDate', 'stock', 'qty', 'etsDate'])

#######################
# Data Analysis / Model
#######################

def fetch_data(q):
    print ("Executing Query:", q)
    result = pd.read_sql(
        sql=q,
        con=conn
    )
    return result

def get_plants():
    '''Returns the list of plants that are stored in the database'''

    plant_query = (
        '''
        SELECT DISTINCT plantCode
        FROM StockReport
        '''
    )
    plants = fetch_data(plant_query)
    plants = list(plants['plantCode'].sort_values(ascending=True))
    return plants

def get_rawMats(plant):
    '''Returns the raw material of selected plant'''

    rawMats_query = (
        '''
        SELECT DISTINCT rmCode
        FROM StockReport
        WHERE plantCode='%s'
        '''%plant
    )
    rawMaterials = fetch_data(rawMats_query)
    rawMaterials = list(rawMaterials['rmCode'].sort_values(ascending=False))
    return rawMaterials

def get_SafeStock(plantCode,rmCode):
    '''Returns the Safe Stock for selected RM and Plant'''

    safeStock_query = (
        '''
        SELECT safeStock
        FROM SafeStock
        WHERE plantCode='%s'
        AND rmCode='%s'
        '''%(plantCode,rmCode)
    )
    safeStockQty = fetch_data(safeStock_query)
    return safeStockQty['safeStock']

#DB query to fetch consumption history
def get_ConsumptionHistory(plantCode, rmCode, reportDate):
    '''Returns consumption history for specified plant, rawMat since specified date'''

    report_query = (
        '''
        SELECT plantCode, rmCode, monthYear, mthlyConsumption
        FROM MthlyConsumption
        WHERE plantCode='%s'
        AND rmCode='%s'
        ORDER BY monthYear ASC
        '''%(plantCode, rmCode)
    )
    consumptionReport = fetch_data(report_query)
    return consumptionReport
        #AND monthYear >='{reportDate}'

#DB query to fetch stock report
def get_StockReport(plantCode, rmCode, reportDate):
    '''Returns stock details for specified plant, rawMat since specified date'''

    report_query = (
        '''
        SELECT plantCode, rmCode, calDate, stock, qty, etsDate
        FROM StockReport
        WHERE plantCode='%s'
        AND rmCode='%s'
        AND calDate >='%s'
        ORDER BY calDate ASC
        '''%(plantCode, rmCode, reportDate)
    )
    stockReport = fetch_data(report_query)
    stockReport =modify_stock_report(stockReport,0,plantCode, rmCode, reportDate)
    return stockReport

def modify_stock_report (results , safeStock,plantCode, rmCode, reportDate):
    # temp hardcode  consumption number
    consumptionReport = get_ConsumptionHistory(plantCode, rmCode, reportDate)
    mnonthly_cosumption = consumptionReport['mthlyConsumption'].iloc[0]
    daily_cosumption = -mnonthly_cosumption/30
    stock = results['stock'].iloc[0]
    start_value = stock
    stop_value = stock +120*daily_cosumption
    #print("series vale" , start_value,stop_value,daily_cosumption)

    data = np.arange(start_value,stop_value,daily_cosumption)
    stock_series = pd.Series(data)
    #print("Stock series is ",len(stock_series))

    date  = results['calDate'].iloc[0]

    plantcode = results['plantCode'].iloc[0]
    rmcode  =results['rmCode'].iloc[0]
    dataframe_new = pd.DataFrame(index=range(120),columns = ['calDate','plantCode', 'rmCode', 'stock', 'qty', 'etsDate'])
    dataframe_new['plantCode'] = plantcode
    dataframe_new['rmCode'] = rmcode
    dataframe_new['stock'].iloc[0] = stock
    dataframe_new['qty'] = pd.to_numeric(results['qty'])
    dataframe_new['etsDate'] = results['etsDate']

    dataframe_new['qty'].fillna(0, inplace=True)
    #print("Report date " ,  pd.to_numeric(dataframe_new['qty']))

    rng = pd.date_range(date, periods=120, freq='D')

   # print("data range datafram", rng.to_series)
    dataframe_new =  dataframe_new.assign(calDate= rng)

    dataframe_new =  dataframe_new.assign(stock= stock_series)

   # print("concated frame", dataframe_new)
    dataframe_new['stock'] =dataframe_new['stock'] +pd.to_numeric(dataframe_new['qty'])
    dataframe_new['stock'] = dataframe_new['stock'].round(2)
    return dataframe_new

def draw_stockReport_graph(results , safeStock):
    dates = results['calDate']
    #safeStocks = [safeStock] * results.shape[0]
    safeStocks = pd.DataFrame([safeStock] * results.shape[0])
    stocks = results['stock']

    # Trace Safe Stock
    trace0 = go.Scatter(
        x = dates,
        y = safeStocks[0],
        name = 'Safe Stock',
        line = dict(
            color = ('rgb(205, 12, 24)'),
            width = 4)
    )
    # Trace Actual Stock
    trace1 = go.Scatter(
        x = dates,
        y = stocks,
        name = 'Actual Stock',
        line = dict(
            color = ('rgb(22, 96, 167)'),
            width = 4,)
    )

    figure = go.Figure(
        data= [ trace0 ,  trace1],
        layout=go.Layout(
            title='Stock Report',
            showlegend=True
        )
    )
    return figure

#########################
# Dashboard Layout / View
#########################

def generate_table(dataframe, max_rows=10):
    '''Given dataframe, return template generated using Dash components
    '''
    return html.Table(
        # Header
        [html.Tr([html.Th(col) for col in dataframe.columns])] +

        # Body
        [html.Tr([
            html.Td(dataframe.iloc[i][col]) for col in dataframe.columns
        ]) for i in range(min(len(dataframe), max_rows))]
    )

def onLoad_plant_options():
    '''Actions to perform upon initial page load'''

    plant_options = (
        [{'label': plantCode, 'value': plantCode}
         for plantCode in get_plants()]
    )
    return plant_options

# Set up Dashboard and create layout
app = dash.Dash(__name__)
app.css.append_css({
    "external_url": "https://codepen.io/chriddyp/pen/bWLwgP.css"
})

app.layout = html.Div([
    html.Div([
        html.Div([
        html.Img(src='/assets/ATGlogo.png',style={'width': '150px'}),
        ], className='three columns'),
        html.Div(className='nine columns'),
    ], className='twelve columns'),
    # Page Header
    html.Div([
        html.H4('Stock Report Viewer')
    ]),

    # Dropdown Grid
    html.Div([
        html.Div([
            # Select Plant Dropdown
            html.Div([
                html.Div('Select Plant', className='three columns'),
                html.Div(dcc.Dropdown(id='plant-selector',
                                      options=onLoad_plant_options()),
                         className='nine columns')
            ]),

            # Select Raw Material Dropdown
            html.Div([
                html.Div('Select Raw Material', className='three columns'),
                html.Div(dcc.Dropdown(id='rawMat-selector'),
                         className='nine columns')
            ]),
        ], className='six columns'),

        # Empty
        # html.Div(className='six columns'),

        html.Div([
             html.P('Select Date'),
             dcc.DatePickerSingle(
        id='reportStartDate',
        min_date_allowed=dt(2001, 8, 1),
        max_date_allowed=dt(2018, 10, 1),
        placeholder='Start Date',
        ),
        ], className='three columns'),

        #Link to trigger /download_excel route
        html.Div(children=[
            html.P(html.A(
                    'Download Data',
                    id='excel-download-link',
                    download="",
                    href="",
                    target="_blank")),
            html.P(html.A(
                    'Download Zip',
                    id='zip-download-link',
                    download="",
                    href="",
                    target="_blank",
		    n_clicks=0)),
            html.P(html.A(['Print PDF'],className="button no-print print")),
        ], className='three columns'),
    ], className='twleve columns'),
    html.Div([
            html.Div([
            # Safe Stock table
            html.Div(children='Safe Stock:',className='three columns'),
            html.Div(id='safe-stock',className='one columns'),
            html.Div(className='two columns'),
            ],className='six columns'),

	    html.Div([
            #Consumption History
            html.Table(id='consumption-history'),
            ],className='six columns'),
    ],className='twelve columns'),

    # Stock Report Grid
    html.Div([

        # Stock Report Table
        html.Div(
            html.Table(id='stock-report'),
            className='six columns'
        ),

        # Stock Summary Table and Graph
        html.Div([
            # graph
            dcc.Graph(id='stock-graph')
            # style={},
        ], className='six columns')
    ]),
])

#############################################
# Interaction Between Components / Controller
#############################################

# Load Seasons in Dropdown
@app.callback(
    Output(component_id='rawMat-selector', component_property='options'),
    [
        Input(component_id='plant-selector', component_property='value')
    ]
)
def populate_rawMaterial_selector(plant):
    rawMaterials = get_rawMats(plant)
    return [
        {'label': rmCode, 'value': rmCode}
        for rmCode in rawMaterials
    ]

# Load Stock Report
@app.callback(
    Output(component_id='stock-report', component_property='children'),
    [
        Input(component_id='plant-selector', component_property='value'),
        Input(component_id='rawMat-selector', component_property='value'),
        Input(component_id='reportStartDate', component_property='date')
    ]
)
def load_stockReport(plantCode, rmCode, reportDate):
    results = get_StockReport(plantCode, rmCode, reportDate)
    return generate_table(results, max_rows=50)

# Load Safe Stock
@app.callback(
    Output(component_id='safe-stock', component_property='children'),
    [
        Input(component_id='plant-selector', component_property='value'),
        Input(component_id='rawMat-selector', component_property='value')
    ]
)
def load_stockReport(plantCode, rmCode):
    safeStock = get_SafeStock(plantCode, rmCode)
    return safeStock

# Load Consumption History
@app.callback(
    Output(component_id='consumption-history', component_property='children'),
    [
        Input(component_id='plant-selector', component_property='value'),
        Input(component_id='rawMat-selector', component_property='value'),
        Input(component_id='reportStartDate', component_property='date')
    ]
)
def load_consumptionHist(plantCode, rmCode, reportDate):
    consumptionHist = get_ConsumptionHistory(plantCode, rmCode, reportDate)
    return generate_table(consumptionHist, max_rows=50)

#Update Stock Point Graph
@app.callback(
    Output(component_id='stock-graph', component_property='figure'),
    [
        Input(component_id='plant-selector', component_property='value'),
        Input(component_id='rawMat-selector', component_property='value'),
        Input(component_id='reportStartDate', component_property='date')
    ]
)
def load_stockReport_graph(plantCode, rmCode, reportDate):
    results = get_StockReport(plantCode, rmCode, reportDate)
    safeStock = get_SafeStock(plantCode, rmCode)

    figure = []
    if len(results) > 0:
        figure = draw_stockReport_graph(results, safeStock)

    return figure

#Callback to download data as csv
#@app.callback(
#    Output(component_id='excel-download-link', component_property='href'),
#    [        Input(component_id='plant-selector', component_property='value'),
#        Input(component_id='rawMat-selector', component_property='value'),
#        Input(component_id='reportStartDate', component_property='date')]
#)
#def download_CSVSingle(plantCode, rmCode, reportDate):
#    stockReport = get_StockReport(plantCode, rmCode, reportDate)
#    globalStockReport = stockReport
#    globalStockReport= globalStockReport.dropna()
#    csv_string = globalStockReport.to_csv(index=False, encoding='utf-8')
#    csv_string = "data:text/csv;charset=utf-8," +p.quote(csv_string)
#    return csv_string

#Callback to download Excel for single RM
@app.callback(
    Output(component_id='excel-download-link', component_property='href'),
    [   Input(component_id='plant-selector', component_property='value'),
        Input(component_id='rawMat-selector', component_property='value'),
        Input(component_id='reportStartDate', component_property='date')]
)
def download_ExcelSingle(plantCode, rmCode, reportDate):
	return url_for('generate_singleRM_report',plantCode=plantCode,rmCode=rmCode,reportDate=reportDate)

#Callback to download ZIP file
@app.callback(
    Output('zip-download-link', 'href'),
    [Input('zip-download-link', 'n_clicks'),
    Input(component_id='plant-selector', component_property='value'),
    Input(component_id='reportStartDate', component_property='date')]
)
def generate_report_url(n_clicks, plantCode, reportDate):
	return url_for('generate_report_url',plantCode=plantCode,reportDate=reportDate)

# Flask routes to serve various files for download

# Serve Excel for single RM
@app.server.route('/dash/singleRMReportDownload')
def generate_singleRM_report():
	plantCode = request.args.get('plantCode')
	rmCode = request.args.get('rmCode')
	reportDate = request.args.get('reportDate')
	stockReport = get_StockReport(plantCode, rmCode, reportDate)
	stockReport = stockReport.dropna()
	rmDownload_filename = str(rmCode) + ".xlsx"
	wsName = str(plantCode) + "_" + str(rmCode)
	absolute_filename = os.path.join(os.getcwd(), rmDownload_filename)
	buf = io.BytesIO()
	#excel_writer = pd.ExcelWriter(absolute_filename, engine="xlsxwriter")
	excel_writer = pd.ExcelWriter(buf, engine="xlsxwriter")
	stockReport.to_excel(excel_writer, sheet_name=wsName)
	excel_writer.save()
	excel_data = buf.getvalue()
	buf.seek(0)
	#return send_file(absolute_filename,
	return send_file(buf,
		attachment_filename = rmDownload_filename, as_attachment = True)
		#mimetype = 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',

# Serve ZIP for All RM
@app.server.route('/dash/allRMZipdownload')
def generate_report_url():
        plantCode = request.args.get('plantCode')
        reportDate = request.args.get('reportDate')
        rawMaterials=get_rawMats(plantCode)

        #Create  excelFile for each RM across the plant
        for rm in rawMaterials:
            rmReport = get_StockReport(plantCode,rm,reportDate)
            rmReport = rmReport.dropna()
            rm_filename = str(rm) + ".xlsx"
            wsName = str(plantCode) + "_" + str(rm)
            relative_filename = os.path.join('excelFiles',rm_filename)
            absolute_filename = os.path.join(os.getcwd(), relative_filename)
            writer = pd.ExcelWriter(absolute_filename, engine='xlsxwriter')
            rmReport.to_excel(writer, sheet_name = wsName)
            writer.save()
        zip_filename = str(plantCode) + ".zip"
        relative_filename = os.path.join('zipFiles',zip_filename)
        absolute_filename = os.path.join(os.getcwd(), relative_filename)
        excelFilePath = os.path.join(os.getcwd(), 'excelFiles')
        base_path = pathlib.Path('./excelFiles')
        with zipfile.ZipFile(relative_filename, mode='w') as z:
            for f_name in base_path.iterdir():
                z.write(f_name)

        return send_file(relative_filename, attachment_filename = zip_filename, as_attachment = True)

external_css = ["https://cdnjs.cloudflare.com/ajax/libs/normalize/7.0.0/normalize.min.css",
                "https://cdnjs.cloudflare.com/ajax/libs/skeleton/2.0.4/skeleton.min.css",
                "//fonts.googleapis.com/css?family=Raleway:400,300,600",
                "https://codepen.io/bcd/pen/KQrXdb.css",
                "https://maxcdn.bootstrapcdn.com/font-awesome/4.7.0/css/font-awesome.min.css"]

for css in external_css:
    app.css.append_css({"external_url": css})

external_js = ["https://code.jquery.com/jquery-3.2.1.min.js",
               "https://codepen.io/bcd/pen/YaXojL.js"]

for js in external_js:
    app.scripts.append_script({"external_url": js})

# start Flask server
if __name__ == '__main__':
    app.run_server(
        debug=True,
        host='0.0.0.0',
        port=8050
    )


#!/usr/bin/env python3

"""
Helpers for being user friendly

"""

import yaspin
from proca.util import log
import click
import configparser
import re
import pprint
from gql.transport.exceptions import TransportQueryError
from asyncio.exceptions import TimeoutError
from aiohttp.client_exceptions import ClientOSError

def explain_error(intent, **fmt_params):
    """
    Decorate function that could throw an error.
    Pass string with %(foo)s strings and their values in fmt_params
    """
    def decor(func):
        def explainer_wrap(*args, **kw):
            try:
                return func(*args, **kw)
            except FileNotFoundError as e:
                fail("ðŸ‘» could not find file: " + intent, fmt_params)
                raise click.Abort(e)
            except IsADirectoryError as e:
                fail("ðŸ¤¨ I expected a file, when "+intent+" but that is a directory", fmt_params)
                raise click.Abort(e)
            except PermissionError as e:
                fail("ðŸ”’ I tried "+intent+" but I do not have permission", fmt_params)
                raise click.Abort(e)
            except configparser.DuplicateSectionError as e:
                fail("ðŸ™ƒ Problem "+intent+" because it already exists", fmt_params)
                raise click.Abort(e)
            except TransportQueryError as e:
                log.debug("GraphQL Errors:", pprint.pformat(e.errors))
                msgs  = ", ".join([api_error_explanation(m) for m in e.errors])
                fail(f"ðŸ˜µ Tried {intent}, but {msgs}")
            except TimeoutError as e:
                log.debug("Timeout error", e)
                fail(f"ðŸ‘Ž Connection to server timed out when {intent}")
            except ClientOSError as e:
                log.debug("Network error", pprint.pformat(e))
                fail(f"ðŸ‘Ž Connection to server broke while {intent}")

        return explainer_wrap
    return decor

def api_error_explanation(msg):
    ex = msg.get('extensions', {})
    if 'path' in msg:
        path = 'at ' + '.'.join(msg['path'])
    else:
        path = ''

    if 'code' in ex:
        if ex['code'] == 'unauthorized':
            return f"you have not provided correct credentials"
        elif ex['code'] == 'bad_arg':
            return f"you have given wrong arguments for {path}: {msg['message']}"
        else:
            return f"server said: {ex['code']} - {msg['message']} {path}"

    return f"{path}: {msg['message']}"

def fail(message, usage=False, exit_code=1):
    if usage:
        error = click.UsageError(message)
    else:
        error = click.ClickException(message)

    error.exit_code = exit_code
    raise error

def make_into_url(url):
    return urL

EMAIL_PATTERN = re.compile(r"(?P<user>^[a-zA-Z][a-zA-Z0-9_.+-]+)@(?P<domain>[a-zA-Z0-9-._]+[a-zA-Z])$")
def validate_email(ctx, param, value):
    if value is None:
        return value

    m = EMAIL_PATTERN.match(value)

    if m is None:
        raise click.BadParameter("user must be an email")

    return value

LOCALE_PATTERN = re.compile(r"^[a-z]{2}(_[A-Z]{2})?$")
def validate_locale(ctx, param, value):
    if value is None:
        return value

    m = LOCALE_PATTERN.match(value)
    if m is None:
        raise click.BadParameter("locale must be language code (en) or language+region code (en_IE)")

    return value

def id_options(fn):
    """
    Adds options to pass different id/name/external ids to the command.
    XXX - add uuid later
    """

    fn = click.option("-n", "--name", type=str, help="short name")(fn)
    fn = click.option("-i", "--id", type=int, help="numerical id")(fn)

    return fn

def guess_identifier(id, name, something):
    pv = len([x for x in [id, name, something] if x is None])
    if pv != 2:
        raise click.UsageError(f"You must specify identifier only in one of three ways: 1) -i 123 2) -n somename 3) 123 or somename as argument.")

    if id is not None:
        return id, None
    if name is not None:
        return None, name
    try:
        return int(something), None
    except ValueError:
        return None, something


from django.db.models.signals import post_save, m2m_changed
from django.dispatch import receiver
from .models import Search, Result, Item, Dataset
from django_celery_results.models import TaskResult
import json

# function to receive signals when new tasks are created and start tasks to do in celery daemon
@receiver([m2m_changed], sender=Search.datasets.through)
def search_save_handler(sender, instance, action, **kwargs):
    if action == "post_add":
        instance.run()

#function to handle new taskresult related to searchs done on local databases
@receiver([post_save], sender=TaskResult)
def search_dataset_result(sender, instance, **kwargs):
    if "search_in_dataset" not in instance.task_name:
        return
    search = Search.objects.filter(query=instance.task_args[0]).latest("query")
    try:
        result = Result.objects.latest(search=search)
    except:
        result = Result(search=search)
        result.save()
    rows = json.loads(instance.result)
    dataset = Dataset.objects.filter(path=instance.task_args[1]).latest("path")
    for row in rows:
        item = Item(row=row, result=result, dataset=dataset)
        item.save()


import argparse
import multiprocessing
import os
import time
import calendar
from pathlib import Path
from typing import List

from . import __version__
from . import sw_logger
from . import slack_wrapper
from . import signature_updater
from .models import (
    signature,
    user,
    workspace,
    post,
    conversation
)

SIGNATURES_PATH = (Path(__file__).parents[2] / 'watchman-signatures').resolve()
OUTPUT_LOGGER: sw_logger.JSONLogger

def load_signatures() -> List[signature.Signature]:
    """ Load signatures from YAML files
    Returns:
        List containing loaded definitions as Signatures objects
    """

    loaded_signatures = []
    try:
        for root, dirs, files in os.walk(SIGNATURES_PATH):
            for sig_file in files:
                sig_path = (Path(root) / sig_file).resolve()
                if sig_path.name.endswith('.yaml'):
                    loaded_def = signature.load_from_yaml(sig_path)
                    for sig in loaded_def:
                        if sig.status == 'enabled' and 'slack_eg' in sig.watchman_apps:
                            loaded_signatures.append(sig)
        return loaded_signatures
    except Exception as e:
        raise e

def search(loaded_signature: signature,
           slack_connection: slack_wrapper.SlackAPI,
           user_list: List[user.User],
           message_list: List[post.Message],
           workspace_list: List[workspace.Workspace],
           files_list: List[post.File],
           scope: str,
           cores: int,
           verbose: bool) -> None:
    """ Uses the signature to call the relevant search functions to find data in messages
    files and drafts. Results are output to stdout logging.
    Args:
        slack_connection: Slack API object
        loaded_signature: Signature object
        user_list: List of User objects from the Enterprise
        message_list: List of Message objects to search through
        workspace_list: List of Workspace objects from the Enterprise
        files_list: List of File objects to search through
        scope: Scope of any results found for logging: e.g. Draft
        cores: Number of CPU cores to use
        verbose: Whether to use verbose logging or not
    """

    if scope == 'messages':
        OUTPUT_LOGGER.log('INFO', f'Searching for posts containing {loaded_signature.name}')
        messages = slack_wrapper.search_message_matches(
            loaded_signature,
            slack_connection,
            user_list,
            message_list,
            workspace_list,
            cores,
            OUTPUT_LOGGER,
            verbose
        )
        if messages:
            for message in messages:
                OUTPUT_LOGGER.log(
                    'NOTIFY',
                    message,
                    scope=scope,
                    severity=loaded_signature.severity,
                    detect_type=loaded_signature.name,
                    notify_type='result'
                )

    if scope == 'files':
        OUTPUT_LOGGER.log('INFO', f'Searching for {loaded_signature.name}')
        files = slack_wrapper.search_file_matches(
            loaded_signature,
            user_list,
            files_list,
            cores,
            OUTPUT_LOGGER
        )
        if files:
            for file in files:
                OUTPUT_LOGGER.log(
                    'NOTIFY',
                    file,
                    scope=scope,
                    severity=loaded_signature.severity,
                    detect_type=loaded_signature.name,
                    notify_type='result'
                )

def core_validation(cores: int) -> int:
    """ Validate the number of cores entered
    Args:
        cores: Number of cores passed to CLI
    Returns:
        Number of cores to use
    """

    if not cores or cores > multiprocessing.cpu_count() or cores > 12:
        if multiprocessing.cpu_count() >= 8:
            return 8
        else:
            return multiprocessing.cpu_count()
    else:
        return cores

def init_logger(logging_type: str, debug: bool) -> sw_logger.JSONLogger or sw_logger.StdoutLogger:
    """ Create a logger object
    Returns:
        Logging object for outputting results
    """

    if not logging_type or logging_type == 'terminal':
        return sw_logger.StdoutLogger(debug=debug)
    else:
        return sw_logger.JSONLogger(debug=debug)

def main():
    global OUTPUT_LOGGER
    try:
        OUTPUT_LOGGER = ''
        parser = argparse.ArgumentParser(description=__version__.__summary__)

        parser.add_argument('--hours', '-hr', dest='hours', type=int,
                            help='How far back to search in whole hours between 1-24. Defaults to 1 if no acceptable '
                                 'value given', required=False)
        parser.add_argument('--minutes', '-m', dest='minutes', type=int,
                            help='How far back to search in whole minutes between 1-60', required=False)
        parser.add_argument('--output', '-o', choices=['json', 'terminal'], dest='logging_type',
                            help='What logging output to use - JSON formatted output, or textual output'
                                 'for reading via terminal. Default is terminal')
        parser.add_argument('--cores', '-c', dest='cores', type=int,
                            help='Number of cores to use between 1-12', required=False)
        parser.add_argument('--version', '-v', action='version',
                            version=f'Slack Watchman for Enterprise Grid: {__version__.__version__}')
        parser.add_argument('--users', '-u', dest='users', action='store_true', help='Return all users')
        parser.add_argument('--workspaces', '-w', dest='workspaces', action='store_true', help='Return all workspaces')
        parser.add_argument('--debug', '-d', dest='debug', action='store_true', help='Turn on debug level logging')
        parser.add_argument('--verbose', '-V', dest='verbose', action='store_true',
                            help='Turn on more verbose output for JSON logging. '
                                 'This includes more fields, but is larger')

        args = parser.parse_args()
        hours = args.hours
        minutes = args.minutes
        cores = args.cores
        users = args.users
        workspaces = args.workspaces
        logging_type = args.logging_type
        debug = args.debug
        verbose = args.verbose

        span = 0
        OUTPUT_LOGGER = init_logger(logging_type, debug)
        if minutes:
            if isinstance(minutes, int) and 1 <= int(minutes) <= 60:
                span = (int(minutes) * 60)
            else:
                minutes = None

        if hours:
            if isinstance(hours, int) and 1 <= int(hours) <= 24:
                span = (span + (hours * 3600))
            else:
                hours = None

        if not minutes and not hours:
            tf = calendar.timegm(time.gmtime()) - 3600
            hours = 1
            minutes = 0
        else:
            tf = calendar.timegm(time.gmtime()) - span

        if not minutes:
            minutes = 0
        if not hours:
            hours = 0

        cores = core_validation(cores)

        slack_con = slack_wrapper.initiate_slack_connection(os.environ.get('SLACK_WATCHMAN_EG_TOKEN'))
        OUTPUT_LOGGER.log('SUCCESS', 'Slack Watchman Enterprise Grid started execution')
        OUTPUT_LOGGER.log('INFO', f'Version: {__version__.__version__}')
        OUTPUT_LOGGER.log('INFO', f'Created by: {__version__.__author__} - {__version__.__email__}')
        OUTPUT_LOGGER.log('INFO', f'{cores} cores in use')
        OUTPUT_LOGGER.log('INFO', 'Downloading signature file updates')
        signature_updater.SignatureUpdater(OUTPUT_LOGGER).update_signatures()
        OUTPUT_LOGGER.log('INFO', 'Importing signatures...')
        signature_list = load_signatures()
        OUTPUT_LOGGER.log('SUCCESS', f'{len(signature_list)} signatures loaded')
        OUTPUT_LOGGER.log('INFO', f'Searching previous {hours} hour(s), {minutes} minutes')
        OUTPUT_LOGGER.log('INFO', 'Enumerating Enterprise information')
        OUTPUT_LOGGER.log('NOTIFY', slack_wrapper.get_enterprise(slack_con), scope='Enterprise',
                          notify_type='enterprise')
        OUTPUT_LOGGER.log('INFO', 'Enumerating Enterprise workspaces')
        workspace_list = slack_wrapper.get_workspaces(slack_con, verbose)
        OUTPUT_LOGGER.log('INFO', f'{len(workspace_list)} workspaces discovered')
        OUTPUT_LOGGER.log('INFO', 'Enumerating Enterprise users')
        user_list = slack_wrapper.get_users(slack_con, workspace_list, verbose)
        OUTPUT_LOGGER.log('INFO', f'{len(user_list)} users discovered')

        if users:
            OUTPUT_LOGGER.log('INFO', 'Outputting Enterprise users')
            for user in user_list:
                OUTPUT_LOGGER.log('NOTIFY', user, detect_type='User', notify_type='user')

        if workspaces:
            OUTPUT_LOGGER.log('INFO', 'Outputting Enterprise workspaces')
            for workspace in workspace_list:
                OUTPUT_LOGGER.log('NOTIFY', workspace, detect_type='Workspace', notify_type='workspace')

        OUTPUT_LOGGER.log('INFO', 'Enumerating files')
        file_list = slack_wrapper.get_all_files(slack_con, cores=cores, verbose=verbose, timeframe=tf)
        OUTPUT_LOGGER.log('INFO', f'{len(file_list)} files discovered')

        OUTPUT_LOGGER.log('INFO', 'Enumerating messages')
        message_list = slack_wrapper.get_all_messages(slack_con, cores=cores, timeframe=tf)
        OUTPUT_LOGGER.log('INFO', f'{len(message_list)} messages discovered')

        for sig in signature_list:
            for scope in sig.scope:
                search(
                    sig,
                    slack_con,
                    user_list,
                    message_list,
                    workspace_list,
                    file_list,
                    scope,
                    cores,
                    verbose
                )

        OUTPUT_LOGGER.log('INFO', 'Enumerating draft messages')
        draft_list = slack_wrapper.get_all_drafts(
            slack_con,
            workspace_list,
            cores=cores,
            verbose=verbose,
            timeframe=tf
        )
        OUTPUT_LOGGER.log('INFO', f'{len(draft_list)} drafts discovered')
        for sig in signature_list:
            if 'drafts' in sig.scope:
                OUTPUT_LOGGER.log('INFO', f'Searching for drafts containing {sig.name}')
                drafts = slack_wrapper.search_draft_matches(
                    slack_con,
                    sig,
                    draft_list,
                    user_list,
                    OUTPUT_LOGGER,
                    verbose,
                    tf
                )
                if drafts:
                    for draft in drafts:
                        OUTPUT_LOGGER.log(
                            'NOTIFY',
                            draft,
                            scope='Draft',
                            severity=sig.severity,
                            detect_type=sig.name,
                            notify_type='result'
                        )

        OUTPUT_LOGGER.log('SUCCESS', 'Slack Watchman Enterprise Grid finished execution')

    except Exception as e:
        OUTPUT_LOGGER.log('CRITICAL', e)

if __name__ == '__main__':
    main()


import numpy as np
import cv2

class Line:
    def __init__(self):
        # was the line detected in the last iteration?
        self.detected = False
        # polynomial coefficients for the most recent fit
        self.current_fit = np.array([False, False, False])
        # x values for detected line pixels
        self.allx = []
        # y values for detected line pixels
        self.ally = []
        # meters per pixel in x dimension
        self.xm_per_pix = 3.7 / 700
        # meters per pixel in y dimension
        self.ym_per_pix = 1 / 30
        # previous fit real for smoothing distance from center calculation
        self.prev_polfit = None

    def get_line_points(self):
        '''
        Returns line's points
        '''
        points = np.dstack((self.allx, self.ally))[0]

        return points

    def draw_line(self, img):
        '''
        Draws line on given image
        '''
        cv2.polylines(img, np.int32([self.get_line_points()]), False, (255, 150, 0), 1)

    def measure_curvature_pixels(self, image):
        '''
        Calculates the curvature of polynomial functions in pixels.
        '''
        # Define y-value where we want radius of curvature
        # We'll choose the maximum y-value, corresponding to the bottom of the image
        y_eval = image.shape[1]

        # Calculation of R_curve (radius of curvature)
        curverad = ((1 + (2 * self.current_fit[0] * y_eval + self.current_fit[1]) ** 2) ** 1.5) / np.absolute(
            2 * self.current_fit[0])

        return curverad

    def measure_curvature_real(self, image):
        '''
        Calculates the curvature of polynomial functions in meters.
        '''
        y_eval = image.shape[1]

        polfit_real = np.polyfit(self.ally * self.ym_per_pix, self.allx * self.xm_per_pix, 2)

        # Calculation of R_curve (radius of curvature)
        curverad = ((1 + (
                    2 * polfit_real[0] * y_eval * self.ym_per_pix + polfit_real[1]) ** 2) ** 1.5) / np.absolute(
            2 * polfit_real[0])

        return curverad

    def measure_distance_real(self, image):
        '''
        Calculates distance to the line
        '''
        polfit = self.current_fit

        if self.prev_polfit is None:
            self.prev_polfit = self.current_fit
        else:
            polfit = self.smooth(self.prev_polfit, self.current_fit)

        f = np.poly1d(polfit)
        middle = image.shape[1] / 2

        return abs(middle - f(image.shape[1])) * self.xm_per_pix

    def smooth(self, prev, curr, coeficient=0.4):
        '''
         exponential smoothing
        :param prev: old value
        :param curr: new value
        :param coeficient: smoothing coef.
        :return:
        '''
        return curr * coeficient + prev * (1 - coeficient)


import logging

class Node:
  def __init__(self, name, value=0, children=None):
    self.logger = logging.getLogger(self.__class__.__qualname__)
    if children is None:
      children = {}
    self.name = name
    self.value = int(value)
    self.children = children

  def add(self, frames, value):
    self.value += int(value)

    if (frames and len(frames) > 0):
      head = frames[0]
      child = self.children[head] if (self.children and head in self.children) else None
      if (child is None):
        child = Node(head)
        self.children[head] = child
      frames[0:1] = []
      child.add(frames, value)

  def serialize(self):
    res = {'name': self.name, 'value': self.value}

    children = []

    for key in self.children:
      children.append(self.children[key].serialize())

    if (len(children) > 0):
      res['children'] = children

    return res