import os
import json
import shutil
import tempfile
import warnings
import subprocess
import numpy as np
import pandas as pd
from astropy.time import Time

from ..utils import writeToADES
from .backend import Backend

FINDORB_CONFIG = {
    "config_file" : os.path.join(os.path.dirname(__file__), "data", "environ.dat"),
}
ADES_KWARGS = {
    "mjd_scale" : "utc",
    "seconds_precision" : 9,
    "columns_precision" : {
        "ra" : 16,
        "dec" : 16,
        "mag" : 2,
        "rmsMag" : 2,
        "rmsTime" : 2
    },
    "observatory_code" : "",
    "submitter" : "",
    "telescope_design" : "",
    "telescope_aperture" : "",
    "telescope_detector" :  "",
    "observers" : [""],
    "measurers" : [""],
    "observatory_name" : "",
    "submitter_institution" : None,
    "telescope_name" : None,
    "telescope_fratio" : None,
    "comment" : None
}

class FINDORB(Backend):

    def __init__(self, **kwargs):

        # Make sure only the correct kwargs
        # are passed to the constructor
        allowed_kwargs = FINDORB_CONFIG.keys()
        for k in kwargs:
            if k not in allowed_kwargs:
                raise ValueError()

        # If an allowed kwarg is missing, add the
        # default
        for k in allowed_kwargs:
            if k not in kwargs:
                kwargs[k] = FINDORB_CONFIG[k]

        super().__init__(name="FindOrb", **kwargs)

        return

    def _writeTimes(self, file_name, times, time_scale):
        """
        Write a time file that find_orb can use to propagate orbits
        and generate ephemerides.

        Parameters
        ----------
        file_name : str
            Path of desired file.
        times :
            Times to write into the file
        time_scale : str
            Time scale to which to ouput to file.
        """
        with open(file_name, "w") as f:
            f.write("OPTION {}\n".format(time_scale.upper()))
            f.write("OPTION ASCENDING\n")
            np.savetxt(
                f,
                times.jd,
                newline="\n",
                fmt="%.16f"
            )
            f.close()
        return

    def _setWorkEnvironment(self, temp_dir):
        # Copy files from user's find_orb directory located in the
        # user's home directory to the .find_orb directory inside
        # the temporary directory
        # Try to ignore files that are created by find_orb when processing
        # observations, also ignore the DE 430 ephemeris file and the asteroid
        # ephemeris file if present
        shutil.copytree(
            os.path.expanduser("~/.find_orb"),
            os.path.join(temp_dir, ".find_orb"),
            ignore=shutil.ignore_patterns(
                "elements.txt",
                "guide.txt",
                "mpc_fmt.txt",
                "observe.txt",
                "residual.txt",
                "total.json",
                "vectors.txt",
                "covar.txt",
                "debug.txt",
                "elements.json",
                "elem_short.json",
                "combined.json",
                "linux_p1550p2650.430t",
                "asteroid_ephemeris.txt"
            )
        )

        # Create a symlink for the DE 430 file that was not copied
        os.symlink(
            os.path.expanduser("~/.find_orb/linux_p1550p2650.430t"),
            os.path.join(temp_dir, ".find_orb/linux_p1550p2650.430t")
        )

        # If the asteroid ephemeris file is present, also create a symlink for this file
        asteroid_ephem_file = os.path.expanduser("~/.find_orb/asteroid_ephemeris.txt")
        if os.path.exists(asteroid_ephem_file):
            os.symlink(
                asteroid_ephem_file,
                os.path.join(temp_dir, ".find_orb/asteroid_ephemeris.txt")
            )

        # Create a copy of the environment and set the HOME directory to the
        # temporary directory
        env = os.environ.copy()
        env["HOME"] = temp_dir
        return env

    def _propagateOrbits(self, orbits, t1, out_dir=None):
        """
        Propagate orbits to t1.

        Parameters
        ----------
        orbits : `~thor.orbits.orbits.Orbits`
            Orbits to propagate.
        t1 : `~astropy.time.core.Time`
            Times to which to propagate each orbit.
        out_dir : str, optional
            Save input and output files to this directory. Will create
            a sub directory called propagation inside this directory.

        Returns
        -------
        propagated : `~pandas.DataFrame`
            Orbits propagated to t1 (sorted by orbit_ids and t1).
        process_returns : list
            List of subprocess CompletedProcess objects.
        """
        propagated_dfs = []
        process_returns = []
        with tempfile.TemporaryDirectory() as temp_dir:

            # Set this environment's home directory to the temporary directory
            # to prevent bugs with multiple processes writing to the ~/.find_orb/
            # directory
            env = self._setWorkEnvironment(temp_dir)
            out_dir_ = os.path.join(temp_dir, "propagation")
            os.makedirs(out_dir_, exist_ok=True)

            # Write the desired times out to a file that find_orb understands
            times_in_file = os.path.join(temp_dir, "propagation", "times_prop.in")
            self._writeTimes(
                times_in_file,
                t1.tt,
                "tt"
            )

            for i in range(orbits.num_orbits):

                # If you give fo a string for a numbered object it will typically append brackets
                # automatically which makes retrieving the object's orbit a little more tedious so by making sure
                # the object ID is not numeric we can work around that.
                orbit_id_i = f"o{i:07}"
                orbit_id_path = "_".join(orbits.ids[i].astype(str).split(" "))
                out_dir_i = os.path.join(out_dir_, orbit_id_path)
                os.makedirs(out_dir_i, exist_ok=True)
                vectors_txt = os.path.join(out_dir_i, "vectors.txt")

                # Format the state vector into the type understood by find_orb
                fo_orbit = "-v{},{}".format(
                    orbits.epochs[i].tt.jd,
                    ",".join(orbits.cartesian[i].astype("str"))
                )

                # Run fo
                call = [
                    "fo",
                    "-o",
                    orbit_id_i,
                    fo_orbit,
                    "-E",
                    "0",
                    "-C",
                    "Sun",
                    "-O",
                    out_dir_i,
                    "-e",
                    vectors_txt,
                    "-D",
                    self.config_file,
                    f"EPHEM_STEP_SIZE=t{times_in_file}"
                ]
                ret = subprocess.run(
                    call,
                    shell=False,
                    env=env,
                    cwd=temp_dir,
                    check=False,
                    capture_output=True,
                )
                process_returns.append(ret)

                if (ret.returncode != 0):
                    warning = (
                        "fo returned a non-zero error code.\n"
                        "Command: \n"
                        " ".join(call) + "\n"
                        f"{ret.stdout.decode('utf-8')}" + "\n"
                        f"{ret.stderr.decode('utf-8')}" + "\n"
                    )
                    warnings.warn(warning)

                if (os.path.exists(vectors_txt)):
                    df = pd.read_csv(
                        vectors_txt,
                        header=0,
                        delim_whitespace=True,
                        names=["jd_tt", "x", "y", "z", "vx", "vy", "vz"]
                    )
                    df["orbit_id"] = [i for _ in range(len(df))]
                else:
                    df = pd.DataFrame(
                        columns=["orbit_id", "jd_tt", "x", "y", "z", "vx", "vy", "vz"]
                    )

                propagated_dfs.append(df)

                if out_dir is not None:

                    # Write a bash script to the temporary directory
                    script_out = os.path.join(out_dir_, f"run_{orbit_id_path}.sh")
                    for i in [9, 11, 13, 14]:
                        call[i] = f'"{call[i]}"'
                    call_out = " ".join(call).replace(
                        out_dir_ + "/",
                        ""
                    )
                    call_out = call_out.replace(
                        os.path.dirname(self.config_file) + "/",
                        ""
                    )
                    with open(script_out, "w") as file_out:
                        file_out.write("#!/bin/sh\n")
                        file_out.write(call_out)

                    # Copy temporary directory to the user defined output directory
                    os.makedirs(out_dir, exist_ok=True)
                    shutil.copytree(
                        temp_dir,
                        out_dir,
                        ignore=shutil.ignore_patterns(
                            ".find_orb",
                        ),
                        dirs_exist_ok=True
                    )

            # Save the configuration file to the user defined output directory (do
            # it only once as its the same configuration for each call)
            if out_dir is not None:
                config_out = os.path.join(out_dir, "propagation", "environ.dat")
                shutil.copy(self.config_file, config_out)

            propagated = pd.concat(propagated_dfs, ignore_index=True)
            propagated["mjd_tdb"] = Time(
                propagated["jd_tt"].values,
                scale="tt",
                format="jd"
            ).tdb.mjd
            propagated = propagated[["orbit_id", "mjd_tdb", "x", "y", "z", "vx", "vy", "vz"]]

            if orbits.ids is not None:
                propagated["orbit_id"] = orbits.ids[propagated["orbit_id"].values]

        return propagated, process_returns

    def _generateEphemeris(self, orbits, observers, out_dir=None):
        """
        Generate ephemerides for each orbit and observer.

        Parameters
        ----------
        orbits : `~thor.orbits.orbits.Orbits`
            Orbits to propagate.
        observers : dict or `~pandas.DataFrame`
            A dictionary with observatory codes as keys and observation_times (`~astropy.time.core.Time`) as values.
        out_dir : str, optional
            Save input and output files to this directory. Will create
            a sub directory called ephemeris inside this directory.

        Returns
        -------
        ephemeris : `~pandas.DataFrame`
            Ephemerides for each orbit and observer.
        process_returns : list
            List of subprocess CompletedProcess objects.
        """
        ephemeris_dfs = []
        process_returns = []
        with tempfile.TemporaryDirectory() as temp_dir:

            # Set this environment's home directory to the temporary directory
            # to prevent bugs with multiple processes writing to the ~/.find_orb/
            # directory
            env = self._setWorkEnvironment(temp_dir)

            for observatory_code, observation_times in observers.items():

                out_dir_ = os.path.join(temp_dir, "ephemeris", observatory_code)
                os.makedirs(out_dir_, exist_ok=True)

                # Write the desired times out to a file that find_orb understands
                times_in_file = os.path.join(out_dir_, "times_eph.in")
                self._writeTimes(
                    times_in_file,
                    observation_times.tt,
                    "tt"
                )

                # Certain observatories are dealt with differently, so appropriately
                # define the columns
                if (observatory_code == "500"):
                    columns = [
                        "jd_utc", "RA_deg", "Dec_deg", "delta_au", "r_au",
                        "elong", "ph_ang", "mag",
                        "'/hr", "PA", "rvel",
                        "lon", "lat", "altitude_km"
                    ]
                else:
                    columns = [
                        "jd_utc", "RA_deg", "Dec_deg", "delta_au", "r_au",
                        "elong", "ph_ang", "mag",
                        "RA_'/hr", "dec_'/hr",
                        "alt",  "az",  "rvel",
                        "lon", "lat", "altitude_km"
                    ]

                # For each orbit calculate their ephemerides
                for i in range(orbits.num_orbits):

                    # If you give fo a string for a numbered object it will typically append brackets
                    # automatically which makes retrieving the object's orbit a little more tedious so by making sure
                    # the object ID is not numeric we can work around that.
                    orbit_id_i = f"o{i:07}"
                    orbit_id_path = "_".join(orbits.ids[i].astype(str).split(" "))
                    out_dir_i = os.path.join(out_dir_, orbit_id_path)
                    os.makedirs(out_dir_i, exist_ok=True)
                    ephemeris_txt = os.path.join(out_dir_i, "ephemeris.txt")

                    # Format the state vector into the type understood by find_orb
                    fo_orbit = "-v{},{}".format(
                        orbits.epochs[i].tt.jd,
                        ",".join(orbits.cartesian[i].astype("str"))
                    )
                    if orbits.H is not None:
                        fo_orbit += ",H={}".format(orbits.H[i])

                    # Call fo and generate ephemerides
                    call = [
                        "fo",
                        "-o",
                        orbit_id_i,
                        fo_orbit,
                        "-C",
                        observatory_code,
                        "-O",
                        out_dir_i,
                        "-e",
                        ephemeris_txt,
                        "-D",
                        self.config_file,
                        f"EPHEM_STEP_SIZE=t{times_in_file}",
                        f"JSON_EPHEM_NAME={os.path.join(out_dir_i, 'eph.json')}",
                        f"JSON_ELEMENTS_NAME={os.path.join(out_dir_i, 'ele.json')}",
                        f"JSON_SHORT_ELEMENTS={os.path.join(out_dir_i, 'short.json')}",
                        f"JSON_COMBINED_NAME={os.path.join(out_dir_i, 'com.json')}"
                    ]

                    ret = subprocess.run(
                        call,
                        shell=False,
                        env=env,
                        cwd=temp_dir,
                        check=False,
                        capture_output=True
                    )
                    process_returns.append(ret)

                    if (ret.returncode != 0):
                        warning = (
                            "fo returned a non-zero error code.\n"
                            "Command: \n"
                            " ".join(call) + "\n"
                            f"{ret.stdout.decode('utf-8')}" + "\n"
                            f"{ret.stderr.decode('utf-8')}" + "\n"
                        )
                        warnings.warn(warning)

                    if (os.path.exists(ephemeris_txt)):
                        ephemeris = pd.read_csv(
                            ephemeris_txt,
                            header=0,
                            delim_whitespace=True,
                            names=columns,
                            float_precision="round_trip"

                        )
                        ephemeris["orbit_id"] = [i for _ in range(len(ephemeris))]
                        ephemeris["observatory_code"] = [observatory_code for _ in range(len(ephemeris))]

                    else:
                        ephemeris = pd.DataFrame(
                            columns=[["orbit_id", "observatory_code"] + columns]
                        )

                    ephemeris_dfs.append(ephemeris)

                    if out_dir is not None:
                        # Write a bash script to the temporary directory
                        script_out = os.path.join(out_dir_, f"run_{orbit_id_path}.sh")
                        for i in [7, 9, 11, 12, 13, 14, 15, 16]:
                            call[i] = f'"{call[i]}"'
                        call_out = " ".join(call).replace(
                            out_dir_ + "/",
                            ""
                        )
                        call_out = call_out.replace(
                            os.path.dirname(self.config_file) + "/",
                            ""
                        )
                        with open(script_out, "w") as file_out:
                            file_out.write("#!/bin/sh\n")
                            file_out.write(call_out)

            if out_dir is not None:
                # Save the configuration file to the temporary directory (do
                # it only once as its the same configuration for each call)
                # This file will be copied with the temporary directory
                config_out = os.path.join(out_dir_, "environ.dat")
                shutil.copy(self.config_file, config_out)

                # Copy temporary directory to the user defined output directory
                os.makedirs(out_dir, exist_ok=True)
                shutil.copytree(
                    temp_dir,
                    out_dir,
                    ignore=shutil.ignore_patterns(
                        ".find_orb",
                        "eph_json.txt"
                    ),
                    dirs_exist_ok=True
                )

        # Combine ephemeris data frames and sort by orbit ID,
        # observatory code and observation time, then reset the
        # index
        ephemeris = pd.concat(ephemeris_dfs)
        ephemeris.sort_values(
            by=["orbit_id", "observatory_code", "jd_utc"],
            inplace=True,
            ignore_index=True
        )

        # Extract observation times and convert them to MJDs
        times = Time(
            ephemeris["jd_utc"].values,
            format="jd",
            scale="utc"
        )
        ephemeris["mjd_utc"] = times.utc.mjd
        ephemeris.drop(
            columns=["jd_utc"],
            inplace=True
        )

        # Sort the columns into a more user friendly order
        cols = ['orbit_id', 'observatory_code', 'mjd_utc'] + [col for col in ephemeris.columns if col not in ['orbit_id', 'observatory_code', 'mjd_utc']]
        ephemeris = ephemeris[cols]

        # If orbits have their IDs defined replace the orbit IDs
        if orbits.ids is not None:
            ephemeris["orbit_id"] = orbits.ids[ephemeris["orbit_id"].values]

        return ephemeris, process_returns

    def _orbitDetermination(self, observations, out_dir=None, ades_kwargs=ADES_KWARGS):
        ids = []
        epochs = []
        orbits = []
        covariances = []
        residual_dfs = []
        process_returns = []

        # Find_Orb accepts ADES files as inputs for orbit determination so
        # lets convert THOR-like observations into essentially dummy
        # ADES files that find_orb can process
        _observations = observations.copy()
        _observations.rename(
            columns={
                "mjd_utc" : "mjd",
                "RA_deg" : "ra",
                "Dec_deg" : "dec",
                "RA_sigma_deg" : "rmsRA",
                "Dec_sigma_deg" : "rmsDec",
                "mag_sigma" : "rmsMag",
                "mjd_sigma_seconds" : "rmsTime",
                "filter" : "band",
                "observatory_code" : "stn",
            },
            inplace=True
        )
        _observations["rmsRA"] = _observations["rmsRA"].values * 3600
        _observations["rmsDec"] = _observations["rmsDec"].values * 3600
        _observations.sort_values(
            by=["mjd", "stn"],
            inplace=True
        )

        id_present = False
        id_col = None
        if "permID" in _observations.columns.values:
            id_present = True
            id_col = "permID"
        if "provID" in _observations.columns.values:
            id_present = True
            id_col = "provID"
        if "trkSub" in _observations.columns.values:
            id_present = True
            id_col = "trkSub"
        if not id_present:
            _observations["trkSub"] = _observations["orbit_id"]
            id_col = "trkSub"

        if "mag" not in _observations.columns:
            _observations.loc[:, "mag"] = 20.0
        if "band" not in _observations.columns:
            _observations.loc[:, "band"] = "V"
        if "mode" not in _observations.columns:
            _observations.loc[:, "mode"] = "CCD"
        if "astCat" not in _observations.columns:
            _observations.loc[:, "astCat"] = "None"

        for i, orbit_id in enumerate(_observations[id_col].unique()):

            with tempfile.TemporaryDirectory() as temp_dir:

                # Set this environment's home directory to the temporary directory
                # to prevent bugs with multiple processes writing to the ~/.find_orb/
                # directory
                env = self._setWorkEnvironment(temp_dir)

                # If you give fo a string for a numbered object it will typically append brackets
                # automatically which makes retrieving the object's orbit a little more tedious so by making sure
                # the object ID is not numeric we can work around that.
                orbit_id_short = f"o{i:07d}"
                if "orbit_id" in _observations.columns:
                    orbit_id_i = "_".join(orbit_id.split(" "))
                else:
                    orbit_id_i = orbit_id_short

                out_dir_ = os.path.join(temp_dir, "orbit_determination")
                out_dir_i = os.path.join(out_dir_, orbit_id_i)
                os.makedirs(out_dir_i, exist_ok=True)

                observations_file = os.path.join(out_dir_, f"observations_{orbit_id_i}.psv")

                mask = _observations[id_col].isin([orbit_id])
                object_observations = _observations[mask].copy()
                object_observations.loc[:, id_col] = orbit_id_short
                object_observations.reset_index(inplace=True, drop=True)

                writeToADES(
                    object_observations,
                    observations_file,
                    **ades_kwargs
                )

                last_observation = Time(
                    object_observations["mjd"].max(),
                    scale="utc",
                    format="mjd"
                )
                call = [
                    "fo",
                    observations_file,
                    "-O",
                    out_dir_i,
                    f"-tEjd{last_observation.tt.jd}",
                    "-j",
                    "-D",
                    self.config_file,
                ]

                ret = subprocess.run(
                    call,
                    shell=False,
                    env=env,
                    cwd=temp_dir,
                    check=False,
                    capture_output=True
                )
                process_returns.append(ret)

                if (ret.returncode != 0):
                    warning = (
                        "fo returned a non-zero error code.\n"
                        "Command: \n"
                        " ".join(call) + "\n"
                        f"{ret.stdout.decode('utf-8')}" + "\n"
                        f"{ret.stderr.decode('utf-8')}" + "\n"
                    )
                    warnings.warn(warning)

                covar_json = os.path.join(out_dir_i, "covar.json")
                if (os.path.exists(covar_json)) and ret.returncode == 0:
                    with open(covar_json) as f:
                        covar_data = json.load(f)
                        epoch = covar_data["epoch"]
                        state = np.array(covar_data["state_vect"])
                        covariance_matrix = np.array(covar_data["covar"])
                else:
                    epoch = np.NaN
                    state = np.zeros(6) * np.NaN
                    covariance_matrix = np.zeros((6,6)) * np.NaN

                total_json = os.path.join(out_dir_i, "total.json")
                if (os.path.exists(total_json)) and ret.returncode == 0:
                    with open(total_json) as f:
                        data = json.load(f)
                        residuals = pd.DataFrame(
                            data["objects"][orbit_id_short]["observations"]["residuals"]
                        )
                        residuals.sort_values(
                            by=["JD", "obscode"],
                            inplace=True
                        )
                        residuals = object_observations[["obs_id"]].join(residuals)

                    residual_dfs.append(residuals)

                ids.append(orbit_id)
                epochs.append(epoch)
                orbits.append(state)
                covariances.append(covariance_matrix)

                if out_dir is not None:
                    # Write a bash script to the temporary directory
                    script_out = os.path.join(out_dir_, f"run_{orbit_id_i}.sh")
                    for i in [1,3,7]:
                        call[i] = f'"{call[i]}"'
                    call_out = " ".join(call).replace(
                        out_dir_ + "/",
                        ""
                    )
                    call_out = call_out.replace(
                        os.path.dirname(self.config_file) + "/",
                        ""
                    )
                    with open(script_out, "w") as file_out:
                        file_out.write("#!/bin/sh\n")
                        file_out.write(call_out)

                    # Copy temporary directory to the user defined output directory
                    os.makedirs(out_dir, exist_ok=True)
                    shutil.copytree(
                        temp_dir,
                        out_dir,
                        ignore=shutil.ignore_patterns(
                            ".find_orb",
                            "eph_json.txt"
                        ),
                        dirs_exist_ok=True
                    )

        if out_dir is not None:
            # Save the configuration file to the user defined output directory (do
            # it only once as its the same configuration for each call)
            config_out = os.path.join(out_dir, "orbit_determination", "environ.dat")
            shutil.copy(self.config_file, config_out)

        orbits = np.vstack(orbits)

        od_orbits = pd.DataFrame({
            "orbit_id" : ids,
            "jd_tt" : epochs,
            "x" : orbits[:, 0],
            "y" : orbits[:, 1],
            "z" : orbits[:, 2],
            "vx" : orbits[:, 3],
            "vy" : orbits[:, 4],
            "vz" : orbits[:, 5],
            "covariance" : covariances
        })

        od_orbits["mjd_tdb"] = np.NaN
        od_orbits.loc[~od_orbits["jd_tt"].isna(), "mjd_tdb"] = Time(
            od_orbits[~od_orbits["jd_tt"].isna()]["jd_tt"].values,
            scale="tt",
            format="jd"
        ).tdb.mjd

        od_orbits = od_orbits[[
            "orbit_id",
            "mjd_tdb",
            "x",
            "y",
            "z",
            "vx",
            "vy",
            "vz",
            "covariance"
        ]]
        od_orbit_members = pd.concat(residual_dfs, ignore_index=True)

        return od_orbits, od_orbit_members, process_returns


"""Definitions for the `AllTimes` class."""
from collections import OrderedDict

import numpy as np
from mosfit.modules.arrays.array import Array
from mosfit.utils import frequency_unit

# Important: Only define one ``Module`` class per file.

class AllTimes(Array):
    """Generate all times for which observations will be constructed.

    Create lists of observation times that associated with real observations
    and interpolations/extrapolations if such flags are passed to MOSFiT.
    """

    def __init__(self, **kwargs):
        """Initialize module."""
        super(AllTimes, self).__init__(**kwargs)
        self._obs_keys = [
            'times', 'bands', 'telescopes', 'systems', 'instruments', 'modes',
            'bandsets', 'frequencies', 'u_frequencies', 'zeropoints',
            'measures', 'observed'
        ]
        self._okeys = [i for i in self._obs_keys if i not in ['observed']]
        for key in self._obs_keys:
            setattr(self, '_' + key, [])

    def process(self, **kwargs):
        """Process module."""
        old_observations = tuple(
            getattr(self, '_' + x) for x in self._obs_keys)
        if (kwargs.get('root', 'output') == 'output'
                and 'extra_times' in kwargs):
            obslist = (
                list(
                    zip(*([kwargs.get(k) for k in self._okeys] +
                          [[True for x in range(len(kwargs['times']))]]))) +
                list(
                    zip(*([kwargs.get('extra_' + k) for k in self._okeys] +
                          [[False
                            for x in range(len(kwargs['extra_times']))]]))))
            obslist.sort()

            self._all_observations = np.concatenate(
                [np.atleast_2d(np.array(x, dtype=object)) for x in obslist],
                axis=0).T
            for ki, key in enumerate(self._obs_keys):
                setattr(self, '_' + key, self._all_observations[ki])
        else:
            for key in list(
                    set(self._obs_keys) - set(['frequencies', 'observed'])):
                setattr(self, '_' + key, kwargs[key])
            self._frequencies = np.array([
                x / frequency_unit(y) if x is not None else None
                for x, y in zip(kwargs['frequencies'], kwargs['u_frequencies'])
            ])
            self._observed = np.full_like(kwargs['times'], True, dtype=bool)
            self._all_observations = tuple(
                getattr(self, '_' + x) for x in self._obs_keys)

        outputs = OrderedDict(
            [('all_' + x, getattr(self, '_' + x))
             for x in list(set(self._obs_keys) - set(['observed']))])
        if any(not np.array_equal(x, y)
               for x, y in zip(old_observations, self._all_observations)):
            self._all_band_indices = np.array(
                [(self._photometry.find_band_index(
                    b, telescope=t, instrument=i, mode=m, bandset=bs, system=s)
                  if f is None else -1) for ti, b, t, s, i, m, bs, f, uf, zps,
                 msr, o in zip(*self._all_observations)])
            self._zps = np.array(
                [x[9] for x in zip(*self._all_observations)])
            self._measures = np.array(
                [x[10] for x in zip(*self._all_observations)])
            self._observation_types = np.array([
                'magcount' if
                ('countrate' in msr and 'magnitude' in msr and zp is not None
                 and self._model._fitter._prefer_fluxes) else
                self._photometry._band_kinds[bi] if bi >= 0 else 'fluxdensity'
                for bi, zp, msr in zip(self._all_band_indices,
                                       self._zps, self._measures)
            ],
                                               dtype=object)
        outputs['all_band_indices'] = self._all_band_indices
        outputs['observation_types'] = self._observation_types
        outputs['observed'] = np.array(self._observed, dtype=bool)
        return outputs

    def receive_requests(self, **requests):
        """Receive requests from other ``Module`` objects."""
        self._photometry = requests.get('photometry', None)


#!/usr/bin/env python
#A program to classify bases as methylated or unmethylated based on long-range signals using the output from nanopolish
#Alexa McIntyre, 2016-2018

import sys
assert sys.version_info >= (3,0), 'please use python3'
import pickle
from collections import defaultdict
import numpy as np
#import time
import os
import glob
import math
import multiprocessing
import warnings
from Bio import SeqIO

from extract_contexts import *
from train_model import train_classifier,pos2label
from read_qual import extract_read_quality
from load_mCaller_data import tsv2matrix

warnings.filterwarnings("ignore", module="sklearn")

def distribute_threads(positions_list,motif,tsvname,read2qual,refname,num_refs,base,mod,nprocs,nvariables,train,modelfile,skip_thresh,qual_thresh,classifier,training_tsv,plot_training):
    """ distributes list of genomic positions across processes then adds resulting signals to matrix if training"""
    outdir='/'.join(tsvname.split('/')[:-1])
    if len(outdir) > 1:
        outdir = outdir + '/'
    print(outdir)
    if not train:
      tsv_output = '.'.join(tsvname.split('.')[:-1])+'.diffs.'+str(nvariables)
      training_pos_dict = None
    else:
      tsv_output = '.'.join(tsvname.split('.')[:-1])+'.diffs.'+str(nvariables)+'.train'
      if training_tsv:
        signal_mat, context_array = tsv2matrix(training_tsv,base)
      else:
          training_pos_dict = pos2label(positions_list)

    print('%d contigs' % num_refs)
    print('%d threads' %nprocs)
    #sys.exit(0)

    if not training_tsv:
        bytesize = os.path.getsize(tsvname)

    if nprocs > 1 and not training_tsv:
        nprocs_allocated = 0
        procs = []
        out_q = multiprocessing.Queue()
        def worker(out_q,tsvname,fastaname,read2qual,nvariables,skip_thresh,qual_thresh,modelfile,classifier,startline,endline,train,training_pos_dict,base=base,motif=None,positions_list=None):
            outtup = extract_features(tsvname,fastaname,read2qual,nvariables,skip_thresh,qual_thresh,modelfile,classifier,startline,endline=endline,train=train,pos_label=training_pos_dict,base=base,motif=motif,positions_list=positions_list)
            out_q.put(outtup)

    if nprocs == 1 and not training_tsv:
        if not train:
            extract_features(tsvname,refname,read2qual,nvariables,skip_thresh,qual_thresh,modelfile,classifier,0,endline=bytesize,train=train,pos_label=training_pos_dict,base=base,motif=motif,positions_list=positions_list)
        else:
            signal_mat, context_array = extract_features(tsvname,refname,read2qual,nvariables,skip_thresh,qual_thresh,modelfile,classifier,0,endline=bytesize,train=train,pos_label=training_pos_dict,base=base,motif=motif,positions_list=positions_list)

    elif nprocs > 1 and not training_tsv:
        chunksize = int(math.ceil(bytesize/float(nprocs)))

        for i in range(nprocs):
            p = multiprocessing.Process(
                    target=worker,
                    args=(out_q,tsvname,refname,read2qual,nvariables,skip_thresh,qual_thresh,modelfile,classifier,chunksize*i,chunksize*(i+1),train,training_pos_dict,base,motif,positions_list))
            procs.append(p)
            p.start()

        if train:
            # Collect all results into a signal matrix and an array of labels
            signal_mat = {}
            context_array = {}
            for i,proc in enumerate(procs):
                tmp_signal_mat,tmp_contexts = out_q.get()
                for twobase in tmp_signal_mat:
                    if twobase not in signal_mat:
                        signal_mat[twobase] = {}
                        context_array[twobase] = {}
                    for label in tmp_signal_mat[twobase]:
                        if label not in signal_mat[twobase]:
                            signal_mat[twobase][label] = []
                            context_array[twobase][label] = []
                        signal_mat[twobase][label].extend(tmp_signal_mat[twobase][label])
                        context_array[twobase][label].extend(tmp_contexts[twobase][label])

        # Wait for all worker processes to finish
        for p in procs:
            p.join ()

    print('Finished extracting signals')
    if not training_tsv:
        tsv_out = tsv_output+'.unsorted.tmp'
        tmpfis = glob.glob(".".join(tsvname.split(".")[:-1])+"*.tmp[0-9]*")
        if nprocs > 1:
            print('Merging files...')
            try:
                os.remove(tsv_output)
            except OSError:
                pass
            for tmpfi in tmpfis:
                os.system('cat ' + tmpfi + ' >> '+ tsv_out )
                os.remove(tmpfi)
            os.system('sort -n -k2 ' + tsv_out + '| uniq > ' + tsv_output)
            os.remove(tsv_out)
        else:
            os.rename(tmpfis[0],tsv_output)

    if train:
       print('Training...')
       #assert len(signal_array) > 5, 'insufficient data aligned to labeled positions for training'
       train_classifier(signal_mat,context_array,modelfile,classifier,plot_training)
       print('Finished training')

def main():
    #parse command line options
    from argparse import ArgumentParser

    parser = ArgumentParser(description='Classify bases as methylated or unmethylated',prog='mCaller')
    all_or_some = parser.add_mutually_exclusive_group(required=True)
    all_or_some.add_argument('-p','--positions',type=str,required=False, help='file with a list of positions at which to classify bases (must be formatted as space- or tab-separated file with chromosome, position, strand, and label if training)')
    all_or_some.add_argument('-m','--motif',type=str,required=False, help='classify every base of type --base in the motif specified instead (can be single one-mer)')
    parser.add_argument('-r','--reference',type=str,required=True,help='fasta file with reference aligned to')
    parser.add_argument('-e','--tsv',type=str,required=True,help='tsv file with nanopolish event alignment')
    parser.add_argument('-f','--fastq',type=str,required=True,help='fastq file with nanopore reads')
    parser.add_argument('-t','--threads',type=int,required=False,help='specify number of processes (default = 1)',default=1)
    parser.add_argument('-b','--base',type=str,required=False,help='bases to classify as methylated or unmethylated (A or C, default A)',default='A')
    parser.add_argument('-n','--num_variables',type=int,required=False,help='change the length of the context used to classify (default of 6 variables corresponds to 11-mer context (6*2-1))',default=6)
    parser.add_argument('--train',action='store_true',required=False,help='train a new model (requires labels in positions file)',default=False)
    parser.add_argument('--training_tsv',type=str,required=False,help='mCaller output file for training')
    parser.add_argument('-d','--modelfile',type=str,required=False,help='model file name')
    parser.add_argument('-s','--skip_thresh',type=int,required=False,help='number of skips to allow within an observation (default 0)',default=0)
    parser.add_argument('-q','--qual_thresh',type=float,required=False,help='quality threshold for reads (default none)',default=0)
    parser.add_argument('-c','--classifier',type=str,required=False,help='use alternative classifier: options = NN (default), RF, LR, or NBC (non-default may significantly increase runtime)',default='NN')
    parser.add_argument('--plot_training',action='store_true',required=False,help='plot probabilities distributions for training positions (requires labels in positions file and --train)',default=False)
    #parser.add_argument('--plot_violin',action='store_true',required=False,help='train a new model (requires labels in positions file)',default=False)
    parser.add_argument('-v','--version',action='version',help='print version',version='%(prog)s v1.0')
    args = parser.parse_args()

    if args.base == 'A':
        mod = 'm6A'
    elif args.base == 'C':
        mod = 'm5C' #TODO: test m4C
    else:
        print('classification only available for A or C bases so far')
        sys.exit(0)

    if not args.modelfile:
        modelfile = os.path.dirname(os.path.realpath(sys.argv[0])) + '/model_'+args.classifier+'_'+str(args.num_variables)+'_'+mod+'.pkl'
    else:
        modelfile = args.modelfile

    if not args.train:
        assert os.path.isfile(modelfile), 'model file not found at '+modelfile

    if args.motif and len(args.motif) == 1:
        base = args.motif
    else:
        base = args.base

    assert (args.skip_thresh < args.num_variables/2), 'too many skips with only '+str(args.num_variables)+' variables - try < half'

    assert os.path.isfile(args.fastq), 'fastq file not found at '+args.fastq
    read2qual = extract_read_quality(args.fastq)

    if not args.training_tsv:
        training_tsv = None
    else:
        training_tsv = args.training_tsv

    try:
        num_refs = 0
        for ref in SeqIO.parse(args.reference,"fasta"):
            num_refs+=1
    except IOError:
        print('reference file missing')
        sys.exit(0)

    #distribute to multiple threads for main computations
    distribute_threads(args.positions,args.motif,args.tsv,read2qual,args.reference,num_refs,base,mod,args.threads,args.num_variables,
        args.train,modelfile,args.skip_thresh,args.qual_thresh,args.classifier,training_tsv,args.plot_training)

if __name__ == "__main__":
    main()


"""This is the entry point into Q|R. It takes user inputs,
   and then constructs all of the objects
   needed to carry out the quantum refinement.

   - fmodel (crystallographic information)
  - calculator (composite object)
  - restraints_manager (computes energy and gradients using either qm codes or
    cctbx (standard))
  - geometry_restraints_manager (analyses geometry e.g. bond RMSDs)
  - weights (scale factors needed to scale up or down data versus restraints
    contributions)
  Then we process them by the refinement/optimization engine, driver.py:

   driver.refine(params   = params,
               fmodel     = fmodel,
               calculator = calculator_manager,
               results    = results)
   results_manager (store all reportable information, and write it out as a log,
   and also write our final pdb structure.)
   """
from __future__ import division

import os
import sys
import time
import pickle

from numpy.lib.function_base import select
import mmtbx.command_line
import mmtbx.f_model
import mmtbx.utils
from libtbx.utils import Sorry
from libtbx import easy_pickle
from libtbx import group_args
from libtbx.utils import null_out
from mmtbx import monomer_library
import mmtbx.monomer_library.pdb_interpretation
import mmtbx.restraints
from fragment import fragments
import calculator
import driver
import restraints
import cluster_restraints
import results
from qrefine.super_cell import expand
import mmtbx.model.statistics
from libtbx import Auto
from ase.io import read as ase_io_read
from cctbx import maptbx

master_params_str ="""

max_atoms = 15000
  .type = int
  .help = maximum number of atoms
debug = False
  .type = bool
  .help = flag to control verbosity of output for debugging problematic code.
cluster{
  fast_interaction = True
    .type = bool
    .help = enable/disable C++ clustering
  charge_cutoff = 8.0
    .type = float
    .help = distance for point charge cutoff
  clustering = True
    .type = bool
    .help = enable/disable clustering
  charge_embedding = False
    .type = bool
    .help = point charge embedding
  two_buffers = False
    .type = bool
    .help = two-buffers are used when gradients for the whole system do not match the joined gradients from fragments \
            when only as single buffer was used to surround each clusters.
  maxnum_residues_in_cluster = 15
    .type = int
    .help = maximum number of residues in a cluster
  select_within_radius = 10
    .type = int
    .help = supersphere expansion radius
  clustering_method = gnc  *bcc
    .type = choice(multi=False)
    .help = type of clustering algorithm
  altloc_method = average *subtract
    .type = choice(multi=False)
    .help = two strategies on how to join energies from multiple energy and \
            gradient calculations are performed for alternate locations.\
            altloc_method=average does not work! (known issue).
  g_scan  = 10 15 20
    .type = str
    .help = sequence of numbers specifying maxnum_residues_in_cluster for gradient convergence test (mode=gtest), treat as string!
  g_mode = None
    .type = int
    .help = manual control over gradient test loops (0=expansion ,1=standard, 2=standard + point-charges, 3=two buffer 4=two_buffer+ point-charges)
  save_clusters = True
    .type = bool
    .help = save currently used fragments and clusters to disk as PDBs.
  bond_with_altloc = True
    .type = bool
    .help = Disable bond_with_altloc check
  re_calculate_rmsd_tolerance = 0.5
    .type = float
    .help = Re-calculate clusters once the model shifted by more than \
            re_calculate_rmsd_tolerance from initial
}

restraints = cctbx *qm
  .type = choice(multi=False)
  .help = Choice of restraints: cctbx (fast) or any of available QM engines.
expansion = False
  .type = bool
  .help = Expand input model into super-sphere
quantum {
  engine_name = *mopac ani torchani terachem turbomole pyscf orca gaussian xtb
    .type = choice(multi=False)
    .help = choose the QM program
  basis = Auto
    .type = str
    .help = pre-defined defaults
  charge= None
    .type = int
    .help = The formal charge of the entire molecule
  method = Auto
    .type = str
    .help = Defaults to HF for all but MOPAC (PM7), xTB (GFN2) and TorchANI (ani-1x_8x)
  memory = None
    .type = str
    .help = memory for the QM program
  nproc = None
    .type = int
    .help = number of parallel processes for the QM program
  qm_addon = gcp dftd3 gcp-d3
    .type = choice(multi=False)
    .help = allows additional calculations of the gCP and/or DFT-D3 corrections using their stand-alone programs
  qm_addon_method = None
    .type = str
    .help = specifies flags for the qm_addon. See manual for details.
}

refine {
  dry_run=False
    .type = bool
    .help = do not perform calculations, only setup steps
  sf_algorithm = *direct fft
    .type = choice(multi=False)
    .help = algorithm used to compute structure factors, either the direct method or fast fourier transform.
  refinement_target_name = *ml ls_wunit_k1
    .type = choice
    .help = Data target type: least-squares or maximum-likelihood.
  mode = opt *refine gtest
    .type = choice(multi=False)
    .help = choose between refinement, geometry optimization or gradient test
  number_of_macro_cycles=1
    .type = int
    .help = number of macro cycles used in the refinement procedure
  number_of_weight_search_cycles=50
    .type = int
    .help = Number of attempts to find optimal weight
  number_of_refine_cycles=5
    .type = int
    .help = maximum number of refinement cycles
  number_of_micro_cycles=50
    .type = int
    .help = maximum number of micro cycles used in refinement
  data_weight=None
    .type = float
    .help = Allows to specify data/restraints weight (by-pass weight search)
  choose_best_use_r_work = False
    .type = bool
    .help = Use Rwork to make a decision about best weight (otherwise use Rfree)
  skip_weight_search = False
    .type = bool
    .help = Calculate weight based on gradient norms (skip weight optimization)
  adjust_restraints_weight_scale_value = 2
    .type = float
  max_iterations_weight = 50
    .type = int
    .help = Max number of trial refinement iterations for weight search
  max_iterations_refine = 50
    .type = int
    .help = Max number of actual refinement iterations
  use_ase_lbfgs = False
    .type = bool
    .help = used for debugging the lbfgs minimizer from cctbx.
  line_search = True
    .type = bool
    .help = flag to use a line search in minimizer.
  stpmax = 0.2
    .type = float
    .help = maximum step length, empirically we find 3 for cctbx, but 0.2 is better for QM methods.
  gradient_only = False
    .type = bool
    .help = use the gradient only line search according to JA Snyman 2005.
  update_all_scales = True
    .type = bool
    .help = Update bulk-solvent and overall scales
  refine_sites = True
    .type = bool
    .help = only refine the cartesian coordinates of the molecular system.
  refine_adp = False
    .type = bool
    .help = adp refinement are not currently supported.
  restraints_weight_scale = 1.0
    .type = float
    .help = Scale factor for restraints term
  shake_sites = False
    .type = bool
    .help = Randomize coordinates prior weight calculation
  use_convergence_test = True
    .type = bool
    .help = Check if refinement converged (for earlier termination)
  max_bond_rmsd = 0.03
    .type = float
    .help = Max bond RMSD to accept optimized weight
  max_r_work_r_free_gap = 5.0
    .type = float
    .help = Max difference between Rfree and Rwork to accept optimized weight
  r_tolerance = 0.001
    .type = float
    .help = R factor delta beteen refinement cycles to determine convergence
  rmsd_tolerance = 0.01
    .type = float
    .help = maximum acceptable tolerance for the rmsd.
  opt_log = False
    .type = bool
    .help = additional output of the L-BFGS optimizer
  pre_opt = False
    .type = bool
    .help = pre-optimization using steepest decent (SD) and conjugate gradient (CG) techniques w/o line search
  pre_opt_stpmax = 0.1
    .type = float
    .help = step size
  pre_opt_iter= 10
    .type = int
    .help = max. iterations for pre-optimizer
  pre_opt_switch = 2
    .type = int
    .help = max. iterations before switching from SD to CG
  pre_opt_gconv = 3000
    .type = float
    .help = gradient norm convergence threshold for pre-optimizer
  minimizer = *lbfgs lbfgsb
    .type = choice(multi=False)
    .help = Choice of minimizer to use. LBFGS-B requires target and gradients.\
            LBFGS can use either gradients only or both, target and gradients.
}

parallel {
  method = *multiprocessing slurm pbs sge lsf threading
    .type = choice(multi=False)
    .help = type of parallel mode and efficient method of processes on the \
            current computer. The others are queueing protocols with the \
            expection of threading which is not a safe choice.
  nproc = None
    .type = int
    .help = Number of processes to use
  qsub_command = None
    .type = str
    .help = Specific command to use on the queue system
}

output_file_name_prefix = None
  .type = str
  .help = Output file name prefix
output_folder_name = "pdb"
  .type = str
  .help = Output folder name
rst_file = None
  .type = str
  .help = Restart file to use for determining location in run. Loads previous \
          results of weight calculations.
dump_gradients=None
  .type = str
  .help = used for debugging gradients when clustering.
"""

def get_master_phil():
  return mmtbx.command_line.generate_master_phil_with_inputs(
    phil_string=master_params_str)

def show_cc(map_data, xray_structure, log=None):
  import mmtbx.maps.mtriage
  from mmtbx.maps.correlation import five_cc
  xrs = xray_structure
  xrs.scattering_type_registry(table="electron")
  d99 = maptbx.d99(
    map = map_data, crystal_symmetry = xrs.crystal_symmetry()).result.d99
  if(log is not None): print >> log, "Resolution of map is: %6.4f" %d99
  result = five_cc(
    map               = map_data,
    xray_structure    = xrs,
    d_min             = d99,
    compute_cc_box    = False,
    compute_cc_mask   = True,
    compute_cc_peaks  = False,
    compute_cc_volume = False).result
  if(log is not None):
    print >> log, "Map-model correlation coefficient (CC)"
    print >> log, "  CC_mask  : %6.4f" %result.cc_mask
  return result.cc_mask

def create_fmodel(cmdline, log):
  fmodel = mmtbx.f_model.manager(
    f_obs          = cmdline.f_obs,
    r_free_flags   = cmdline.r_free_flags,
    xray_structure = cmdline.xray_structure,
    target_name    = cmdline.params.refine.refinement_target_name)
  if(cmdline.params.refine.update_all_scales):
    fmodel.update_all_scales(remove_outliers=False)
    fmodel.show(show_header=False, show_approx=False)
  print >> log, "Initial r_work=%6.4f r_free=%6.4f" % (fmodel.r_work(), fmodel.r_free())
  log.flush()
  return fmodel

def process_model_file(pdb_file_name, cif_objects=None, crystal_symmetry=None):
  import iotbx.pdb
  params = mmtbx.model.manager.get_default_pdb_interpretation_params()
  params.pdb_interpretation.use_neutron_distances = True
  params.pdb_interpretation.restraints_library.cdl = False
  params.pdb_interpretation.sort_atoms = False
  pdb_inp = iotbx.pdb.input(file_name=pdb_file_name)
  model = mmtbx.model.manager(
    model_input       = pdb_inp,
    crystal_symmetry  = crystal_symmetry,
    restraint_objects = cif_objects,
    log               = null_out())
  model.process(make_restraints=True, grm_normalization=True,
    pdb_interpretation_params = params)
  return model

def create_fragment_manager(
      model,
      params,
      file_name = os.path.join("ase","tmp_ase.pdb")):
  if(not params.cluster.clustering): return None
  return fragments(
    cif_objects                = model._restraint_objects,
    working_folder             = os.path.split(file_name)[0]+ "/",
    clustering_method          = params.cluster.clustering_method,
    altloc_method              = params.cluster.altloc_method,
    maxnum_residues_in_cluster = params.cluster.maxnum_residues_in_cluster,
    charge_embedding           = params.cluster.charge_embedding,
    two_buffers                = params.cluster.two_buffers,
    pdb_hierarchy              = model.get_hierarchy(),
    qm_engine_name             = params.quantum.engine_name,
    crystal_symmetry           = model.crystal_symmetry(),
    debug                      = params.debug,
    fast_interaction           = params.cluster.fast_interaction,
    charge_cutoff              = params.cluster.charge_cutoff,
    save_clusters              = params.cluster.save_clusters,
    select_within_radius       = params.cluster.select_within_radius,
    bond_with_altloc_flag      = params.cluster.bond_with_altloc)

def create_restraints_manager(params, model):
  restraints_source = restraints.restraints(params = params, model = model)
  if(params.expansion):
    if(model.model.altlocs_present()):
      return restraints.from_altlocs(
        restraints_source = restraints_source,
        pdb_hierarchy     = model.model.get_hierarchy(),
        crystal_symmetry  = model.crystal_symmetry,
        method            = params.cluster.altloc_method)
    else:
      return restraints.from_expansion(
        params           = params,
        restraints_source = restraints_source,
        pdb_hierarchy     = model.model.get_hierarchy(),
        crystal_symmetry  = model.crystal_symmetry)
  else:
    if(params.cluster.clustering):
      fragment_manager = create_fragment_manager(params = params, model = model)
      return cluster_restraints.from_cluster(
        restraints_manager = restraints_source.restraints_manager,
        fragment_manager   = fragment_manager,
        parallel_params    = params.parallel)
    else:
      return restraints_source.restraints_manager

def create_calculator(weights, params, restraints_manager, fmodel=None,
                      model=None):
  if(weights is None):
    weights = calculator.weights(
      shake_sites             = params.refine.shake_sites ,
      restraints_weight       = 1.0,
      data_weight             = params.refine.data_weight,
      restraints_weight_scale = params.refine.restraints_weight_scale)
  if(params.refine.refine_sites):
    if(params.refine.mode == "refine"):
      return calculator.sites(
        fmodel             = fmodel,
        restraints_manager = restraints_manager,
        weights            = weights,
        dump_gradients     = params.dump_gradients)
    else:
      return calculator.sites_opt(
        restraints_manager = restraints_manager,
        model              = model,
        dump_gradients     = params.dump_gradients,
        max_shift          = params.refine.stpmax)
  if(params.refine.refine_adp):
    return calculator.adp(
      fmodel             = fmodel,
      restraints_manager = restraints_manager,
      weights            = weights)

def validate(model, fmodel, params, rst_file, prefix, log):
  # set defaults
  outl = ''
  if params.quantum.engine_name=='mopac':
    if params.quantum.method==Auto:
      params.quantum.method='PM7'
      outl += '  Setting QM method to PM7\n'
    if params.quantum.basis==Auto:
      params.quantum.basis=''
  if params.quantum.engine_name=='xtb':
    if params.quantum.method==Auto:
      params.quantum.method=' --gfn 2 --etemp 500 --acc 0.1 --gbsa h2o'
      print >> log, '  Default method for xtb is %s' % (
          params.quantum.method,
          )
      params.quantum.basis = ''
  else:
    if params.quantum.method==Auto:
      params.quantum.method='HF'
      outl += '  Setting QM method to HF\n'
    if params.quantum.basis==Auto:
      params.quantum.basis='6-31g'
      outl += '  Setting QM basis to 6-31g\n'
  if outl:
    print >> log, '\nSetting QM defaults'
    print >> log, outl

  if params.quantum.engine_name=='mopac':
    if params.quantum.basis:
      print >> log, '  Because engine is %s basis set %s ignored' % (
        params.quantum.engine_name,
        params.quantum.basis,
        )
      params.quantum.basis = ''
    if params.quantum.method=='hf': # default
      print >> log, '  Default method set as PM7'
      params.quantum.method='PM7'

def run_g_test(params, model, weights, start_fmodel, log):
  # needs to be moved! Perhaps also to driver
  import numpy as np
  from fragment import fragment_extracts, write_cluster_and_fragments_pdbs
  from utils.mathbox import get_grad_mad, get_grad_angle

  # determine what kind of buffer to calculate
  g_mode=[]
  if params.cluster.g_mode is None:
    g_mode.append(1)
    if params.cluster.charge_embedding:
      g_mode.append(2)
    if params.cluster.two_buffers:
      g_mode.append(3)
    if params.cluster.two_buffers and params.cluster.charge_embedding:
      g_mode.append(4)
  else:
    g_mode.append(params.cluster.g_mode)

  # reset flags
  params.cluster.clustering=True
  params.cluster.save_clusters=True
  params.cluster.charge_embedding=False
  params.cluster.two_buffers=False
  grad=[]
  idx=0
  idl=[]

  # input for cluster size
  cluster_scan=sorted([int(x) for x in params.cluster.g_scan.split()])
  # cluster_scan=[2,10]

  if g_mode[0]==0:
    print >> log, 'warning: supersphere calculation!'
    params.cluster.clustering=False
    params.expansion=True
    cluster_scan=[0]
    clusters=[]

  n_grad=len(cluster_scan)*len(g_mode)
  print >> log, 'Calculating %3i gradients \n' % (n_grad)
  print >> log, 'Starting loop over different fragment sizes'
  for ig in g_mode:

    print >> log,'loop for g_mode = %i ' % (ig)
    if ig == 2:
      print >> log, 'pc on'
      params.cluster.charge_embedding=True
    if ig == 3:
      print >> log, 'two_buffers on, pc off'
      params.cluster.charge_embedding=False
      params.cluster.two_buffers=True
    if ig == 4:
      print >> log, 'two_buffers on, pc on'
      params.cluster.charge_embedding=True
      params.cluster.two_buffers=True

    for max_cluster in cluster_scan:
      idl.append([ig,max_cluster])
      print >> log, 'g_mode: %s' % (" - ".join(map(str,idl[idx])))
      t0 = time.time()
      print >> log, "~max cluster size ",max_cluster
      params.cluster.maxnum_residues_in_cluster=max_cluster
      fragment_manager = create_fragment_manager(params = params, model = model)
      restraints_manager = create_restraints_manager(params, model)
      rm = restraints_manager
      if(fragment_manager is not None):
        print "time taken for fragments",(time.time() - t0)
        frags=fragment_manager
        print >> log, '~  # clusters  : ',len(frags.clusters)
        print >> log, '~  list of atoms per cluster:'
        print >> log, '~   ',[len(x) for x in frags.cluster_atoms]
        print >> log, '~  list of atoms per fragment:'
        print >> log, '~   ',[len(x) for x in frags.fragment_super_atoms]

        # save fragment data. below works
        # better way is to make a single PDB file with chain IDs
        label="-".join(map(str,idl[idx]))
        write_cluster_and_fragments_pdbs(fragments=fragment_extracts(frags),directory=label)

      calculator_manager = create_calculator(
        weights            = weights,
        fmodel             = start_fmodel,
        model              = model,
        params             = params,
        restraints_manager = rm)
      grad=list(calculator_manager.target_and_gradients())[1]
      print >> log, '~   gnorm',np.linalg.norm(grad)
      print >> log, '~   max_g', max(abs(i) for i in grad), ' min_g',min(abs(i) for i in grad)
      name="-".join(map(str,idl[idx]))
      np.save(name,grad)
      idx+=1
      print >> log, "total time for gradient",(time.time() - t0),'\n\n'

  print >> log, 'ready to run qr.granalyse!'

def run(model, fmodel, map_data, params, rst_file, prefix, log):
  validate(model, fmodel, params, rst_file, prefix, log)
  if(params.cluster.clustering):
    params.refine.gradient_only = True
    print >> log, " params.gradient_only", params.refine.gradient_only
  # RESTART
  if(os.path.isfile(str(rst_file))):
    print >> log, "restart info is loaded from %s" % params.rst_file
    rst_data = easy_pickle.load(params.rst_file)
    fmodel = rst_data["fmodel"]
    results_manager = rst_data["results"]
    results_manager.log = log
    weights = rst_data["weights"]
    geometry_rmsd_manager = rst_data["geometry_rmsd_manager"]
    start_fmodel = rst_data["rst_fmodel"]
    start_ph = model.pdb_hierarchy.deep_copy().adopt_xray_structure(
      start_fmodel.xray_structure)
  else:
    weights = None
    if (model.size() > params.max_atoms):
      raise Sorry("Too many atoms.")
    geometry_rmsd_manager = model.get_restraints_manager()
    #
    if(params.refine.dry_run): return
    #
    r_work, r_free = None, None
    if(fmodel is not None):
      r_work, r_free = fmodel.r_work(), fmodel.r_free()
    results_manager = results.manager(
      r_work                  = r_work,
      r_free                  = r_free,
      model                   = model,
      max_bond_rmsd           = params.refine.max_bond_rmsd,
      max_r_work_r_free_gap   = params.refine.max_r_work_r_free_gap,
      mode                    = params.refine.mode,
      log                     = log,
      restraints_weight_scale = params.refine.restraints_weight_scale)
    if(params.rst_file is None):
      if(params.output_file_name_prefix is None):
        params.output_file_name_prefix = prefix
      if(os.path.exists(params.output_folder_name) is False):
        os.mkdir(params.output_folder_name)
      params.rst_file = os.path.abspath(params.output_folder_name + "/" + \
        params.output_file_name_prefix + ".rst.pickle")
    if os.path.isfile(params.rst_file):
      os.remove(params.rst_file)
    print >> log, "\n***********************************************************"
    print >> log, "restart info will be stored in %s" % params.rst_file
    print >> log, "***********************************************************\n"
    start_fmodel = fmodel
    start_ph = None # is it used anywhere? I don't see where it is used!

  restraints_manager = create_restraints_manager(params, model)

  if(params.refine.mode == "gtest"):
    run_g_test(params=params, log=log, model=model, weights=weights,
               start_fmodel = start_fmodel)
  else:
    if(map_data is not None and params.refine.mode == "refine"):
      model.geometry_statistics(use_hydrogens=False).show()
      show_cc(
        map_data        = map_data,
        xray_structure  = model.get_xray_structure(),
        log             = log)
      O = calculator.sites_real_space(
        model                 = model,
        geometry_rmsd_manager = geometry_rmsd_manager,
        max_bond_rmsd         = params.refine.max_bond_rmsd,
        map_data              = map_data,
        data_weight           = params.refine.data_weight,
        refine_cycles         = params.refine.number_of_refine_cycles,
        skip_weight_search    = params.refine.skip_weight_search,
        stpmax                = params.refine.stpmax,
        gradient_only         = params.refine.gradient_only,
        line_search           = params.refine.line_search,
        restraints_manager    = restraints_manager,
        max_iterations        = params.refine.max_iterations_refine,
        log                   = log)
      model = O.run()
      of = open("real_space_refined.pdb", "w")
      print >> of, model.model_as_pdb(output_cs=True)
      of.close()
      model.geometry_statistics(use_hydrogens=False).show()
      show_cc(
        map_data=map_data,
        xray_structure=model.get_xray_structure(),
        log=log)
      return
    else:
      calculator_manager = create_calculator(
        weights=weights,
        fmodel=start_fmodel,
        model=model,
        params=params,
        restraints_manager=restraints_manager)
      if(params.refine.mode == "refine"):
        driver.refine(
          params                = params,
          fmodel                = fmodel,
          geometry_rmsd_manager = geometry_rmsd_manager,
          calculator            = calculator_manager,
          results               = results_manager)
      else:
        driver.opt(
          params     = params,
          model      = model,
          calculator = calculator_manager,
          results    = results_manager)
      xrs_best = results_manager.finalize(
        input_file_name_prefix  = prefix,
        output_file_name_prefix = params.output_file_name_prefix,
        output_folder_name      = params.output_folder_name,
        use_r_work              = params.refine.choose_best_use_r_work)

if (__name__ == "__main__"):
  t0 = time.time()
  log = sys.stdout
  args = sys.argv[1:]
  print >> log, '_'*80
  print >> log, 'Command line arguments'
  outl = '  '
  for arg in args:
    outl += '"%s"' % arg
  print >> log, '%s\n' % outl
  print >> log, '_'*80
  cmdline = mmtbx.command_line.load_model_and_data(
      args          = args,
      master_phil   = get_master_phil(),
      create_fmodel = False,
      out           = log)
  run(cmdline=cmdline, log = log)
  print >> log, "Time: %6.4f"%(time.time()-t0)


import logging
import logging.config

import configparser as cp
import args

import numpy as np
from keras.preprocessing.sequence import pad_sequences
import pickle

import vectorizer
import load
import sample

#configurations
config = cp.ConfigParser(strict=False)
config.read('defaults.cfg')
#################################################
config['cnn']['filters'] = '50'
config['cnn']['optimizer'] = 'adam'
config['cnn']['lr'] = '0.0001'
config['cnn']['loss'] = 'binary_crossentropy'
config['cnn']['dropout'] = '0'
config['embedding']['length'] = '20'
config['embedding']['limit'] = '1000000'
#################################################

#argparser
args = args.get_args()
'''
>>> args.train
False
'''

#logging
logger = logging.getLogger(__name__)
logging.config.dictConfig({
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {
        'standard': {
            'format': '%(asctime)s [%(levelname)s] %(name)s: %(message)s'
        },
    },
    'handlers': {
        'default': {
            'level':'INFO',
            'formatter': 'standard',
            'class':'logging.StreamHandler',
        },
    },
    'loggers': {
        '': {
            'handlers': ['default'],
            'level': 'INFO',
            'propagate': True
        }
    }
})

[real_val_data,concept_order] = pickle.load(open('gitig_real_val_data.pickle','rb'))
real_val_data.y=np.array(real_val_data.y)

#truncated real eval data
#[real_val_data,concept_order,mention_padded] = pickle.load(open('gitig_real_val_data_truncated.pickle','rb'))
#real_val_data.y=np.array(real_val_data.y)

# word embedding
vector_model, vocabulary, inversed_vocabulary = vectorizer.prepare_embedding_vocab('/home/lhchan/disease_normalization/data/pubmed2018_w2v_400D/pubmed2018_w2v_400D.bin', binary = True, limit = 1000000)
# vector_model, vocabulary, inversed_vocabulary = vectorizer.prepare_embedding_vocab('/home/lenz/disease-normalization/data/embeddings/wvec_50_haodi-li-et-al.bin', binary = True, limit = 1000000)
pretrained = vectorizer.load_pretrained_word_embeddings(vocabulary, vector_model)

# MEDIC dictionary
dictionary = load.Terminology()
# dictionary of entries, key = canonical id, value = named tuple in the form of
#   MEDIC_ENTRY(DiseaseID='MESH:D005671', DiseaseName='Fused Teeth',
#   AllDiseaseIDs=('MESH:D005671',), AllNames=('Fused Teeth', 'Teeth, Fused')
dictionary.loaded = load.load(config['terminology']['dict_file'],'MEDIC')

def concept_obj(conf,dictionary,order=None):
    # concept_ids = [] # list of all concept ids
    # concept_all_ids = [] # list of (lists of all concept ids with alt IDs)
    concept_names = [] # list of all names, same length as concept_ids
    # concept_map = {} # names as keys, ids as concepts

    if order:
        use = order
        logger.info('Re-initialing concept object.')
    else:
        use = dictionary.loaded.keys()

    for k in use:
    # keys not in congruent order! To make them congruent:
    # k,v = zip(*dictionary.loaded.items())
    # k = list(k)
    # k.sort()
        # c_id = dictionary.loaded[k].DiseaseID
        # a_ids = dictionary.loaded[k].AllDiseaseIDs
        for n in dictionary.loaded[k].AllNames:
            concept_names.append(n)

    # tokenization & vectorization of dictionary terms
    import nltk
    concept_tokenize = [nltk.word_tokenize(name) for name in concept_names] # list of list of tokenized names
    concept_vectorize = np.array([[vocabulary.get(text.lower(),1) for text in concept] for concept in concept_tokenize])

    # save the stuff to object
    concept = sample.NewDataSet('concepts')
    #concept.ids = concept_ids
    #concept.all_ids = concept_all_ids
    concept.names = concept_names
    #concept.map = concept_map
    concept.tokenize = concept_tokenize
    concept.vectorize = concept_vectorize
    for corpus in [concept]:
        logger.info('Padding {0}'.format(corpus.info))
        logger.info('Old shape: {0}'.format(corpus.vectorize.shape))
        corpus.padded = pad_sequences(corpus.vectorize, padding='post', maxlen=int(config['embedding']['length']))
        #format of corpus.padded: numpy, mentions, padded
        logger.info('New shape: {0}'.format(corpus.padded.shape))

    return concept

concept = concept_obj(config,dictionary,order=concept_order)

# corpus
corpus_dev = sample.NewDataSet('dev corpus')
corpus_dev.objects = load.load(config['corpus']['development_file'],'NCBI')

for corpus in [corpus_dev]:
    # mention_ids = [] # list of all ids (gold standard for each mention)
    mention_names = [] # list of all names
    # mention_all = [] # list of tuples (mention_text,gold,context,(start,end,docid))

    #sth wrong here that sometimes throw an error
    #import pdb;pdb.set_trace()
    for abstract in corpus.objects:
        for section in abstract.sections: # title and abstract
            for mention in section.mentions:
                # nor_ids = [sample._nor_id(one_id) for one_id in mention.id]
                # mention_ids.append(nor_ids) # append list of ids, usually len(list)=1
                mention_names.append(mention.text)
                # mention_all.append((mention.text,nor_ids,section.text,(mention.start,mention.end,abstract.docid)))

    # tokenization & vectorization of mentions
    import nltk
    mention_tokenize = [nltk.word_tokenize(name) for name in mention_names]
    mention_vectorize = np.array([[vocabulary.get(text.lower(),1) for text in mention] for mention in mention_tokenize])

    # corpus.ids = mention_ids
    corpus.names = mention_names
    # corpus.all = mention_all
    corpus.tokenize = mention_tokenize
    corpus.vectorize = mention_vectorize

for corpus in [corpus_dev]:
    logger.info('Padding {0}'.format(corpus.info))
    logger.info('Old shape: {0}'.format(corpus.vectorize.shape))
    corpus.padded = pad_sequences(corpus.vectorize, padding='post', maxlen=int(config['embedding']['length']))
    #format of corpus.padded: numpy, mentions, padded
    logger.info('New shape: {0}'.format(corpus.padded.shape))

# load the model to be evaluated
from cnn import semantic_similarity_layer
#from keras.models import load_model
#model = load_model('models/20190324-220005.h5',custom_objects={'semantic_similarity_layer': semantic_similarity_layer})
from model_tools import load_model
import cnn, model_tools, callback
from keras.models import Model
from keras.layers import Input, Embedding, Concatenate, Dense, Conv1D, GlobalMaxPooling1D, Flatten
from keras import layers
model_shared_encoder = cnn.build_model_shared_encoder(config,real_val_data,vocabulary,pretrained)
model_shared_encoder.load_weights('models/synpair_continued_68.h5')

model_shared_encoder.summary()

# def forward_pass_speedup_shared_encoder(model,corpus_padded,concept_padded,pretrained):
#     '''
#     Model to speed up forward pass, used in callback for evaluation
#     '''
#     model_mention = _forward_pass_speedup_conv(model,['inp_mentions','embedding_1','conv1d_1','global_max_pooling1d_1'])
#     mentions = model_mention.predict(corpus_padded) # (787, 50)
#     model_candidate = _forward_pass_speedup_conv(model,['inp_candidates','embedding_1','conv1d_1','global_max_pooling1d_2'])
#     candidates = model_candidate.predict(concept_padded) # (67782,50)
#     logger.info('Formatting pooled mentions and candidates...')
#     # from sample import no_cangen_format_x
#     from sample import sped_up_format_x
#     convoluted_input = sped_up_format_x(mentions,candidates)
#     model_sem = _forward_pass_speedup_sem(model,convoluted_input)
#     return convoluted_input, model_sem

# def _forward_pass_speedup_sem(original_model,convoluted_x):
#     layers = ['semantic_similarity_layer_1','dense_1','dense_2']
#     v_sem = original_model.get_layer(layers[0])
#     d1 = original_model.get_layer(layers[1])
#     d2 = original_model.get_layer(layers[2])

#     pooled_mentions = Input(shape=(convoluted_x[0].shape[1],),dtype='float32', name='pooled_mentions')
#     pooled_candidates = Input(shape=(convoluted_x[1].shape[1],),dtype='float32', name='pooled_candidates')
#     sem = semantic_similarity_layer(weights = v_sem.get_weights())([pooled_mentions,pooled_candidates])
#     concatenate_list = [pooled_mentions,pooled_candidates,sem]
#     join_layer = Concatenate()(concatenate_list)
#     hidden_layer = Dense(d1.units, activation=d1.activation,weights=d1.get_weights())(join_layer)
#     prediction_layer = Dense(d2.units, activation=d2.activation,weights=d2.get_weights())(hidden_layer)

#     input_list = [pooled_mentions, pooled_candidates]
#     model_part = Model(inputs=input_list, outputs=prediction_layer)
#     return model_part

# def _forward_pass_speedup_conv(original_model,layers):
#     '''
#     Input:
#     original_model
#     layers: list of layer names, one of the two
#         ['inp_mentions','embedding_1','conv1d_1','global_max_pooling1d_1']
#         ['inp_candidates','embedding_1','conv1d_2','global_max_pooling1d_2']
#     '''
#     terms = original_model.get_layer(layers[0])
#     emb = original_model.get_layer(layers[1])
#     conv = original_model.get_layer(layers[2])

#     new_input_terms = Input(shape=(terms.input_shape[1],),dtype='int32', name='new_input_terms')
#     new_emb = Embedding(emb.input_dim, emb.output_dim, mask_zero=False, trainable=False, weights=emb.get_weights())
#     encoded = new_emb(new_input_terms)
#     new_conv = Conv1D(filters=conv.filters,kernel_size=conv.kernel_size[0],activation=conv.activation,weights=conv.get_weights())(encoded)
#     gl_max_p = GlobalMaxPooling1D()(new_conv)

#     model_part = Model(inputs=new_input_terms, outputs=gl_max_p)
#     return model_part

def _forward_pass_speedup_conv(original_model,layerss):
    '''
    Input:
    original_model
    layerss: list of layer names, one of the two
        ['inp_mentions','embedding_1','drop',conv1d_1','global_max_pooling1d_1']
        ['inp_candidates','embedding_1','drop','conv1d_2','global_max_pooling1d_2']
    '''
    terms = original_model.get_layer(layerss[0])
    emb = original_model.get_layer(layerss[1])
    drop = original_model.get_layer(layerss[2])
    conv = original_model.get_layer(layerss[3])

    new_input_terms = Input(shape=(terms.input_shape[1],),dtype='int32', name='new_input_terms')
    new_emb = Embedding(emb.input_dim, emb.output_dim, mask_zero=False, trainable=False, weights=emb.get_weights())
    new_drop = layers.Dropout(drop.get_config()['rate'])
    encoded = new_drop(new_emb(new_input_terms))
    new_conv = Conv1D(filters=conv.filters,kernel_size=conv.kernel_size[0],activation=conv.activation,weights=conv.get_weights())(encoded)
    gl_max_p = GlobalMaxPooling1D()(new_conv)

    model_part = Model(inputs=new_input_terms, outputs=gl_max_p)
    return model_part

def _forward_pass_speedup_sem(original_model,convoluted_x):
    layers = ['v_sem','hidden_layer','prediction_layer']
    v_sem = original_model.get_layer(layers[0])
    d1 = original_model.get_layer(layers[1])
    d2 = original_model.get_layer(layers[2])

    pooled_mentions = Input(shape=(convoluted_x[0].shape[1],),dtype='float32', name='pooled_mentions')
    pooled_candidates = Input(shape=(convoluted_x[1].shape[1],),dtype='float32', name='pooled_candidates')
    sem = semantic_similarity_layer(weights = v_sem.get_weights())([pooled_mentions,pooled_candidates])
    concatenate_list = [pooled_mentions,pooled_candidates,sem]
    join_layer = Concatenate()(concatenate_list)
    hidden_layer = Dense(d1.units, activation=d1.activation,weights=d1.get_weights())(join_layer)
    prediction_layer = Dense(d2.units, activation=d2.activation,weights=d2.get_weights())(hidden_layer)

    input_list = [pooled_mentions, pooled_candidates]
    model_part = Model(inputs=input_list, outputs=prediction_layer)
    return model_part

def forward_pass_speedup_shared_encoder(model,corpus_padded,concept_padded,pretrained):
    '''
    Model to speed up forward pass, used in callback for evaluation
    '''
    model_mention = _forward_pass_speedup_conv(model,['inp_mentions','embedding_layer','drop','conv1d','global_max_pooling1d_1'])
    mentions = model_mention.predict(corpus_padded) # (787, 50)
    model_candidate = _forward_pass_speedup_conv(model,['inp_candidates','embedding_layer','drop','conv1d','global_max_pooling1d_2'])
    candidates = model_candidate.predict(concept_padded) # (67782,50)
    logger.info('Formatting pooled mentions and candidates...')
    # from sample import no_cangen_format_x
    from sample import sped_up_format_x
    convoluted_input = sped_up_format_x(mentions,candidates)
    model_sem = _forward_pass_speedup_sem(model,convoluted_input)
    return convoluted_input, model_sem

conv_x, sped_up_model = forward_pass_speedup_shared_encoder(model_shared_encoder,corpus_dev.padded,concept.padded,pretrained)
tr_predictions = sped_up_model.predict(conv_x)
acc = callback.evaluate(real_val_data.mentions,tr_predictions,real_val_data.y)
logger.info('Accuracy for the pretrained model on validation set:{0}'.format(acc))


# feature test using threads and thread events rather than gpio.event
# WORKS!!!
# Using queue to communicate which page the user is currently on (relevant as commands are in different classes)
##
# V1 - a window with a button which opens up another window with a video playing
# V2 - the basic interface works with buttons which link to video streams
#    - way found to change button focus
# V3 - added keyboard navigation
# V4 - need to create IR interface - Opi-gpio
# V4-ir - IR Commands feature tests only
# V5 - need ip connection -
# V6 - need to add style (stylesheets)

import OPi.GPIO as GPIO
from time import sleep
from time import perf_counter
import threading
import numpy
from queue import Queue

LED_gpio = 7 # pin name = PA6 wPi = 2 GPIO = 6
ir_gpio = 5 # pin name = SCL.0 wPi = 1 GPIO = 11

GPIO.setmode(GPIO.BOARD)  # set up BOARD BCM SUNXI numbering
#GPIO.setup(LED_gpio, GPIO.OUT)    # set pin as an output (LED)

GPIO.setup(ir_gpio, GPIO.IN)

# ===================================================================================================
# IR controls

#wait_time = 4523982e-9 * 1.5 # 60e-3 # in seconds should be just higher than the longest wait. think longest is about 5ms. noted from looking theough sky remote commands
wait_time = 6e-3 # in seconds should be just higher than the longest wait. think longest is about 5ms. noted from looking theough sky remote commands

ir_keys_dict = { # in nanosecs
    'down' : numpy.array([0.004283, 0.004550, 0.000543, 0.001578, 0.000513, 0.001692, 0.000649, 0.001594, 0.000513, 0.000588, 0.000662, 0.000485, 0.000484, 0.000617, 0.000567, 0.000519, 0.000568, 0.000509, 0.000566, 0.001751, 0.000541, 0.001657, 0.000485, 0.001736, 0.000489, 0.000604, 0.000567, 0.000562, 0.000499, 0.000485, 0.000729, 0.000488, 0.000497, 0.000496, 0.000638, 0.001626, 0.000650, 0.000479, 0.000497, 0.000642, 0.000564, 0.000503, 0.000561, 0.000659, 0.000485, 0.001689, 0.000498, 0.001710, 0.000598, 0.000490, 0.000561, 0.000511, 0.000643, 0.001634, 0.000573, 0.001625, 0.000628, 0.001652, 0.000480, 0.001717, 0.000618, 0.000522, 0.000665, 0.000479, 0.000477, 0.001691, 0.000502]),
    'up' :   numpy.array([0.004288, 0.004426, 0.000647, 0.001535, 0.000654, 0.001583, 0.000607, 0.001568, 0.000646, 0.000572, 0.000502, 0.000524, 0.000615, 0.000494, 0.000654, 0.000489, 0.000597, 0.000466, 0.000613, 0.001711, 0.000514, 0.001665, 0.000554, 0.001641, 0.000568, 0.000521, 0.000586, 0.000543, 0.000509, 0.000579, 0.000655, 0.000484, 0.000583, 0.000415, 0.000721, 0.000511, 0.000581, 0.000489, 0.000569, 0.000487, 0.000673, 0.000543, 0.000506, 0.000489, 0.000644, 0.001623, 0.000598, 0.001611, 0.000504, 0.000730, 0.000498, 0.001597, 0.000655, 0.001596, 0.000595, 0.001588, 0.000651, 0.001574, 0.000528, 0.001732, 0.000564, 0.000562, 0.000509, 0.000509, 0.000670, 0.001537, 0.000645]),
    'right': numpy.array([0.004273, 0.004527, 0.000435, 0.001788, 0.000553, 0.001568, 0.000650, 0.001661, 0.000498, 0.000498, 0.000651, 0.000486, 0.000564, 0.000620, 0.000599, 0.000496, 0.000490, 0.000487, 0.000694, 0.001671, 0.000442, 0.001704, 0.000644, 0.001608, 0.000577, 0.000570, 0.000493, 0.000566, 0.000567, 0.000518, 0.000588, 0.000491, 0.000646, 0.000565, 0.000484, 0.000487, 0.000733, 0.001488, 0.000646, 0.000484, 0.000565, 0.000604, 0.000535, 0.000579, 0.000566, 0.001605, 0.000673, 0.001552, 0.000568, 0.000566, 0.000565, 0.001646, 0.000566, 0.000483, 0.000644, 0.001573, 0.000752, 0.001458, 0.000726, 0.001658, 0.000499, 0.000496, 0.000647, 0.000484, 0.000565, 0.001636, 0.000581]),
    'left' : numpy.array([0.004381, 0.004456, 0.000502, 0.001620, 0.000570, 0.001685, 0.000605, 0.001628, 0.000533, 0.000601, 0.000520, 0.000524, 0.000587, 0.000568, 0.000491, 0.000566, 0.000607, 0.000514, 0.000587, 0.001682, 0.000514, 0.001677, 0.000520, 0.001691, 0.000519, 0.000582, 0.000586, 0.000590, 0.000493, 0.000486, 0.000649, 0.000579, 0.000502, 0.000519, 0.000509, 0.001704, 0.000591, 0.000488, 0.000716, 0.001613, 0.000573, 0.000520, 0.000520, 0.000603, 0.000497, 0.001645, 0.000623, 0.001688, 0.000506, 0.000489, 0.000567, 0.000643, 0.000523, 0.001666, 0.000576, 0.000491, 0.000567, 0.001599, 0.000679, 0.001665, 0.000439, 0.000667, 0.000509, 0.000504, 0.000581, 0.001731, 0.000496]),
    'exit':  numpy.array([0.004383, 0.004593, 0.000486, 0.001635, 0.000487, 0.001646, 0.000673, 0.001685, 0.000485, 0.000515, 0.000626, 0.000412, 0.000731, 0.000481, 0.000486, 0.000638, 0.000596, 0.000499, 0.000509, 0.001728, 0.000500, 0.001665, 0.000569, 0.001675, 0.000502, 0.000604, 0.000670, 0.000485, 0.000487, 0.000664, 0.000557, 0.000446, 0.000531, 0.000612, 0.000504, 0.001715, 0.000587, 0.000544, 0.000524, 0.001671, 0.000637, 0.001564, 0.000644, 0.000514, 0.000428, 0.001779, 0.000538, 0.000513, 0.000650, 0.000484, 0.000484, 0.000560, 0.000654, 0.001689, 0.000547, 0.000430, 0.000743, 0.000507, 0.000501, 0.001644, 0.000584, 0.000485, 0.000570, 0.001711, 0.000567, 0.001685, 0.000510]),
    'enter' :numpy.array([0.004291, 0.004423, 0.000652, 0.001543, 0.000591, 0.001699, 0.000569, 0.001668, 0.000520, 0.000576, 0.000595, 0.000495, 0.000564, 0.000558, 0.000560, 0.000516, 0.000576, 0.000595, 0.000486, 0.001727, 0.000540, 0.001636, 0.000485, 0.001676, 0.000745, 0.000482, 0.000500, 0.000482, 0.000721, 0.000482, 0.000479, 0.000599, 0.000569, 0.000589, 0.000486, 0.000560, 0.000565, 0.000638, 0.000545, 0.000414, 0.000747, 0.001520, 0.000632, 0.000519, 0.000490, 0.001703, 0.000643, 0.001622, 0.000498, 0.000650, 0.000562, 0.001622, 0.000601, 0.001602, 0.000567, 0.001720, 0.000485, 0.000507, 0.000566, 0.001676, 0.000678, 0.000532, 0.000511, 0.000490, 0.000645, 0.001650, 0.000573]),
    'return':numpy.array([0.004415, 0.004415, 0.000515, 0.001658, 0.000700, 0.001538, 0.000517, 0.001694, 0.000580, 0.000528, 0.000570, 0.000516, 0.000569, 0.000560, 0.000482, 0.000662, 0.000592, 0.000494, 0.000583, 0.001670, 0.000505, 0.001686, 0.000589, 0.001619, 0.000646, 0.000419, 0.000578, 0.000511, 0.000570, 0.000564, 0.000538, 0.000598, 0.000591, 0.000490, 0.000585, 0.000569, 0.000486, 0.000565, 0.000670, 0.000511, 0.000574, 0.001635, 0.000483, 0.001727, 0.000590, 0.000491, 0.000645, 0.001560, 0.000591, 0.000578, 0.000500, 0.001681, 0.000608, 0.001649, 0.000492, 0.001758, 0.000523, 0.000494, 0.000668, 0.000491, 0.000482, 0.001730, 0.000489, 0.000503, 0.000732, 0.001629, 0.000505])
}

ir_keys_tolerance_dict = {
    'up' : 0.000187,
    'down' : 0.000309,
    'left' : 0.000237,
    'right' : 0.000253,
    'exit' : 0.000271,
    'enter' : 0.000235,
    'return' : 0.000246
}

ir_keys_length = {
    'up' : len(ir_keys_dict['up']),
    'down' : len(ir_keys_dict['down']),
    'left' : len(ir_keys_dict['left']),
    'right' : len(ir_keys_dict['right']),
    'exit' : len(ir_keys_dict['exit']),
    'enter' : len(ir_keys_dict['enter']),
    'return' : len(ir_keys_dict['return'])
}

class IR_reciever:

    def __init__(self, pin_num, delay_time):
        self.pin_num = pin_num
        self.delay_time = delay_time
        self.q = Queue(maxsize = 1) #create queue for communicating with ir thread. this will allow changing remote functions
        #print(q.qsize())

    def change_screen(self, message):
        self.q.put(message)

    def wait_for_falling_edge(self):

        while thread_flag_event.is_set() : #uses threading event to dictate when the loop should be killed
            if GPIO.input(ir_gpio) == 0:
                self.read_command(0)

            if self.q.full():
                print("queue message detected the message was:")
                mesg = self.q.get()
                #print(mesg)
                if mesg == 'main screen':
                    print('main screen message detected ')
                elif mesg == 'menu screen':
                    print('menu screen message detected ')
                elif mesg == 'video screen':
                    print('video screen message detected ')

    def wait_for_rising_edge(self):
        while thread_flag_event.is_set() :
            if GPIO.input(ir_gpio) == 1:
                self.read_command(1)

    def stop_detection(self):
        self.run_in_background == False

    def read_command(self, start_val):
        print("input detected")
        command = []
        start_of_event = perf_counter()
        IR_timer = start_of_event + wait_time #make it extra long to stop the loop exiting premiturely. value will be reset inside loop

        previousValue = bool(start_val)

        while perf_counter() < IR_timer: #waits until a timer expires. the timer refers to the time after a command has not been passed. this is readjusted instide the code to make sure the entire command is captured
            #print(GPIO.input(ir_gpio)) #dont debug with this, it slows everything down too much and it misses all the inputs
            if GPIO.input(ir_gpio) != previousValue:  # Waits until change in state occurs.
                # Records the current time              maybe change to just record microseconds
                end_of_event = perf_counter() #records the end time of a signal change. the first start_of_event is above the while loop so from here out we need to record just hte$
                # Calculate time in between pulses
                command_length = end_of_event - start_of_event
                #print(command_length)
                start_of_event = perf_counter() #Resets the start time

                command.append(command_length)

                IR_timer = perf_counter() + wait_time #resets counter

                previousValue = not previousValue #value #changes previous value to current value to see if there is a change

        #print("finished event detect, lengh of command = ", len(command), "  command recorded = ")
        #print(command)
        #print("")

        len_command = len(command)

        if min(ir_keys_length.values()) == len_command:
            self.translate_command(command) #compare value

        #if min(ir_keys_length['down']) == len_command: #check if the lengths are the same, all the commands are the same
        #    self.translate_command(command) #compare value

    def translate_command(self, command):
        np_command = numpy.asarray(command, dtype=numpy.float32) # dtype=numpy.uintc)
        #print("checking command translation ")
        #for i in ir_keys_dict: #scroll though the list >> just idea not finished
        #    if numpy.allclose(np_command, ir_keys_dict[i], rtol=1 , atol= ir_keys_tolerance_dict[i]): #checks if the arrays are the same within a tolerance
        #        #check what command the match worked with
        #        #excecute the command
        #        #break for loop

        if numpy.allclose(np_command, ir_keys_dict['up'], rtol=1 , atol= ir_keys_tolerance_dict['up']): #checks if the arrays are the same within a tolerance
            print("up button pressed")
            #self.screen_choice_class.navigation_control('up')
            #self.MainScreen_class.navigation_control('up')

        if numpy.allclose(np_command, ir_keys_dict['down'], rtol=1 , atol= ir_keys_tolerance_dict['down']): #down
            print('down remote button pressed')
            #self.screen_choice_class.navigation_control('down')   ####

        if numpy.allclose(np_command, ir_keys_dict['left'], rtol=1 , atol= ir_keys_tolerance_dict['left']):  #left
            print('left remote button pressed')
            #self.screen_choice_class.navigation_control('left')

        if numpy.allclose(np_command, ir_keys_dict['right'], rtol=1 , atol= ir_keys_tolerance_dict['right']): #right
            print('right remote button pressed')
            #self.screen_choice_class.navigation_control('right')

        if numpy.allclose(np_command, ir_keys_dict['enter'], rtol=1 , atol= ir_keys_tolerance_dict['enter']): #enter button
            print ("enter remote button pressed")
            #self.screen_choice_class.navigation_control('enter')

        if numpy.allclose(np_command, ir_keys_dict['return'], rtol=1 , atol= ir_keys_tolerance_dict['return']): #backup button not needed for main screen
            print ("backup remote button pressed")
            #self.screen_choice_class.navigation_control('return')

        if numpy.allclose(np_command, ir_keys_dict['exit'], rtol=1 , atol= ir_keys_tolerance_dict['exit']): #backup button not needed for main screen
            print ("exit button pressed")
            #self.screen_choice_class.navigation_control('exit')

try:
    gpio_input = IR_reciever(ir_gpio, wait_time)  #initialte the function by inputting the input pin number and the wait time (longest time between individual pulses

    thread_flag_event = threading.Event()
    thread_flag_event.set()
    thread = threading.Thread(target = gpio_input.wait_for_falling_edge) #place this into a thread so that it can run in parallel with the
    thread.start()

    test_number = 0
    print("running program press ctrl+c to stop")
    while 1:
        sleep(10)
        if test_number == 0:
            gpio_input.change_screen('main screen')
            test_number = test_number+1
        elif test_number == 1:
            gpio_input.change_screen('menu screen')
            test_number = test_number+1
        else:
            gpio_input.change_screen('video screen')
            test_number = 0

except KeyboardInterrupt:
    thread_flag_event.clear() #tell the thread to shut down
    print("")
    print("program stopping")
    GPIO.cleanup()