# coding: utf-8
# Copyright (c) Tingzheng Hou.
# Distributed under the terms of the MIT License.

"""
This module implements functions for coordination analysis.
"""

from typing import Dict, List, Tuple, Union, Callable, Optional

import numpy as np
from tqdm.notebook import tqdm
from MDAnalysis import Universe, AtomGroup
from MDAnalysis.core.groups import Atom
from MDAnalysis.analysis.distances import distance_array
from scipy.signal import savgol_filter
from mdgo.util import atom_vec, angle

__author__ = "Tingzheng Hou"
__version__ = "1.0"
__maintainer__ = "Tingzheng Hou"
__email__ = "tingzheng_hou@berkeley.edu"
__date__ = "Feb 9, 2021"

def neighbor_distance(
    nvt_run: Universe,
    center_atom: Atom,
    run_start: int,
    run_end: int,
    species: str,
    select_dict: Dict[str, str],
    distance: float,
) -> Dict[str, np.ndarray]:
    """
    Calculates a dictionary of distances between the ``center_atom`` and neighbor atoms.

    Args:
        nvt_run: An MDAnalysis ``Universe`` containing wrapped trajectory.
        center_atom: The center atom object.
        run_start: Start frame of analysis.
        run_end: End frame of analysis.
        species: The neighbor species in the select_dict.
        select_dict: A dictionary of atom species selection, where each atom species name is a key
            and the corresponding values are the selection language.
        distance: The neighbor cutoff distance.

    Returns:
        A dictionary of distance of neighbor atoms to the ``center_atom``. The keys are atom indexes in string type .
    """
    dist_dict = {}
    time_count = 0
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    species_selection = select_dict.get(species)
    if species_selection is None:
        raise ValueError("Invalid species selection")
    for ts in trj_analysis:
        selection = (
            "(" + species_selection + ") and (around " + str(distance) + " index " + str(center_atom.index) + ")"
        )
        shell = nvt_run.select_atoms(selection, periodic=True)
        for atom in shell.atoms:
            if str(atom.index) not in dist_dict:
                dist_dict[str(atom.index)] = np.full(run_end - run_start, 100.0)
        time_count += 1
    time_count = 0
    for ts in trj_analysis:
        for atom_index, val in dist_dict.items():
            dist = distance_array(ts[center_atom.index], ts[int(atom_index)], ts.dimensions)
            val[time_count] = dist
        time_count += 1
    return dist_dict

def find_nearest(
    trj: Dict[str, np.ndarray],
    time_step: float,
    binding_cutoff: float,
    hopping_cutoff: float,
    smooth: int = 51,
) -> Tuple[List[int], Union[float, np.floating], List[int]]:
    """Using the dictionary of neighbor distance ``trj``, finds the nearest neighbor ``sites`` that the center atom
    binds to, and calculates the ``frequency`` of hopping between each neighbor, and ``steps`` when each binding site
    exhibits the closest distance to the center atom.

    Args:
        trj: A dictionary of distances between center atom and neighbor atoms.
        time_step: The time step of the simulation in ps.
        binding_cutoff: Binding cutoff distance.
        hopping_cutoff: Detaching cutoff distance.
        smooth: The length of the smooth filter window. Default to 51.

    Returns:
        Returns an array of nearest neighbor ``sites`` (unique on each frame),
        the ``frequency`` of hopping between sites, and ``steps`` when each binding site
        exhibits the closest distance to the center atom.
    """
    time_span = len(list(trj.values())[0])
    if smooth > 0:
        for kw in list(trj):
            trj[kw] = savgol_filter(trj.get(kw), smooth, 2)
    site_distance = [100 for _ in range(time_span)]
    sites: List[Union[int, np.integer]] = [0 for _ in range(time_span)]
    start_site = min(trj, key=lambda k: trj[k][0])
    kw_start = trj.get(start_site)
    assert kw_start is not None
    if kw_start[0] < binding_cutoff:
        sites[0] = int(start_site)
        site_distance[0] = kw_start[0]
    else:
        pass
    for time in range(1, time_span):
        if sites[time - 1] == 0:
            old_site_distance = 100
        else:
            old_trj = trj.get(str(sites[time - 1]))
            assert old_trj is not None
            old_site_distance = old_trj[time]
        if old_site_distance > hopping_cutoff:
            new_site = min(trj, key=lambda k: trj[k][time])
            new_trj = trj.get(new_site)
            assert new_trj is not None
            new_site_distance = new_trj[time]
            if new_site_distance > binding_cutoff:
                site_distance[time] = 100
            else:
                sites[time] = int(new_site)
                site_distance[time] = new_site_distance
        else:
            sites[time] = sites[time - 1]
            site_distance[time] = old_site_distance
    sites = [int(i) for i in sites]
    sites_and_distance_array = np.array([[sites[i], site_distance[i]] for i in range(len(sites))])
    steps = []
    closest_step: Optional[int] = 0
    previous_site = sites_and_distance_array[0][0]
    if previous_site == 0:
        closest_step = None
    for i, step in enumerate(sites_and_distance_array):
        site = step[0]
        distance = step[1]
        if site == 0:
            pass
        else:
            if site == previous_site:
                if distance < sites_and_distance_array[closest_step][1]:
                    closest_step = i
                else:
                    pass
            else:
                if closest_step is not None:
                    steps.append(closest_step)
                closest_step = i
                previous_site = site
    if closest_step is not None:
        steps.append(closest_step)
    change = (np.diff([i for i in sites if i != 0]) != 0).sum()
    assert change == len(steps) - 1 or change == len(steps) == 0
    frequency = change / (time_span * time_step)
    return sites, frequency, steps

def find_nearest_free_only(
    trj: Dict[str, np.ndarray],
    time_step: float,
    binding_cutoff: float,
    hopping_cutoff: float,
    smooth: int = 51,
) -> Tuple[List[int], Union[float, np.floating], List[int]]:
    """Using the dictionary of neighbor distance ``trj``, finds the nearest neighbor ``sites`` that the ``center_atom``
    binds to, and calculates the ``frequency`` of hopping between each neighbor, and ``steps`` when each binding site
    exhibits the closest distance to the center atom.
    * Only hopping events with intermediate free state (no binded nearest neighbor) are counted.

    Args:
        trj: A dictionary of distances between center atom and neighbor atoms.
        time_step: The time step of the simulation in ps.
        binding_cutoff: Binding cutoff distance.
        hopping_cutoff: Detaching cutoff distance.
        smooth: The length of the smooth filter window. Default to 51.

    Returns:
        Returns an array of nearest neighbor ``sites`` (unique on each frame),
        the ``frequency`` of hopping between sites, and ``steps`` when each binding site
        exhibits the closest distance to the center atom.
    """
    time_span = len(list(trj.values())[0])
    if smooth > 0:
        for kw in list(trj):
            trj[kw] = savgol_filter(trj.get(kw), smooth, 2)
    site_distance = [100 for _ in range(time_span)]
    sites: List[Union[int, np.integer]] = [0 for _ in range(time_span)]
    start_site = min(trj, key=lambda k: trj[k][0])
    kw_start = trj.get(start_site)
    assert kw_start is not None
    if kw_start[0] < binding_cutoff:
        sites[0] = int(start_site)
        site_distance[0] = kw_start[0]
    else:
        pass
    for time in range(1, time_span):
        if sites[time - 1] == 0:
            old_site_distance = 100
        else:
            old_trj = trj.get(str(sites[time - 1]))
            assert old_trj is not None
            old_site_distance = old_trj[time]
        if old_site_distance > hopping_cutoff:
            new_site = min(trj, key=lambda k: trj[k][time])
            new_trj = trj.get(new_site)
            assert new_trj is not None
            new_site_distance = new_trj[time]
            if new_site_distance > binding_cutoff:
                site_distance[time] = 100
            else:
                sites[time] = int(new_site)
                site_distance[time] = new_site_distance
        else:
            sites[time] = sites[time - 1]
            site_distance[time] = old_site_distance
    sites = [int(i) for i in sites]
    sites_and_distance_array = np.array([[sites[i], site_distance[i]] for i in range(len(sites))])
    steps = []
    closest_step: Optional[int] = 0
    previous_site = sites_and_distance_array[0][0]
    previous_zero = False
    if previous_site == 0:
        closest_step = None
        previous_zero = True
    for i, step in enumerate(sites_and_distance_array):
        site = step[0]
        distance = step[1]
        if site == 0:
            previous_zero = True
        else:
            if site == previous_site:
                if distance < sites_and_distance_array[closest_step][1]:
                    closest_step = i
                else:
                    pass
            elif not previous_zero:
                previous_site = site
                if distance < sites_and_distance_array[closest_step][1]:
                    closest_step = i
                else:
                    pass
            else:
                if closest_step is not None:
                    steps.append(closest_step)
                closest_step = i
                previous_site = site
    if closest_step is not None:
        steps.append(closest_step)
    frequency = (len(steps) - 1) / (time_span * time_step)
    return sites, frequency, steps

def find_in_n_out(
    trj: Dict[str, np.ndarray], binding_cutoff: float, hopping_cutoff: float, smooth: int = 51, cool: int = 20
) -> Tuple[List[int], List[int]]:
    """Finds the frames when the center atom binds with the neighbor (binding) or hopping out (hopping)
    according to the dictionary of neighbor distance.

    Args:
        trj: A dictionary of distances between center atom and neighbor atoms.
        binding_cutoff: Binding cutoff distance.
        hopping_cutoff: Hopping out cutoff distance.
        smooth: The length of the smooth filter window. Default to 51.
        cool: The cool down frames between hopping in and hopping out. Default to 20.

    Returns:
        Two arrays of numberings of frames with hopping in and hopping out event, respectively.
    """
    time_span = len(list(trj.values())[0])
    if smooth > 0:
        for kw in list(trj):
            trj[kw] = savgol_filter(trj.get(kw), smooth, 2)
    site_distance = [100 for _ in range(time_span)]
    sites = [0 for _ in range(time_span)]
    start_site = min(trj, key=lambda k: trj[k][0])
    kw_start = trj.get(start_site)
    assert kw_start is not None
    if kw_start[0] < binding_cutoff:
        sites[0] = int(start_site)
        site_distance[0] = kw_start[0]
    else:
        pass
    for time in range(1, time_span):
        if sites[time - 1] == 0:
            old_site_distance = 100
        else:
            old_trj = trj.get(str(sites[time - 1]))
            assert old_trj is not None
            old_site_distance = old_trj[time]
        if old_site_distance > hopping_cutoff:
            new_site = min(trj, key=lambda k: trj[k][time])
            new_trj = trj.get(new_site)
            assert new_trj is not None
            new_site_distance = new_trj[time]
            if new_site_distance > binding_cutoff:
                site_distance[time] = 100
            else:
                sites[time] = int(new_site)
                site_distance[time] = new_site_distance
        else:
            sites[time] = sites[time - 1]
            site_distance[time] = old_site_distance
    sites = [int(i) for i in sites]

    last = sites[0]
    steps_in: List[int] = []
    steps_out: List[int] = []
    in_cool = cool
    out_cool = cool
    for i, s in enumerate(sites):
        if last == s:
            pass
        elif last == 0:
            in_cool = 0
            steps_in.append(i)
            if out_cool < cool:
                steps_out.pop()
        elif s == 0:
            out_cool = 0
            steps_out.append(i)
            if in_cool < cool:
                steps_in.pop()
        else:
            if cool == 0:
                steps_out.append(i - 1)
                steps_in.append(i)
        last = s
        in_cool += 1
        out_cool += 1
    return steps_in, steps_out

def check_contiguous_steps(
    nvt_run: Universe,
    center_atom: Atom,
    distance_dict: Dict[str, float],
    select_dict: Dict[str, str],
    run_start: int,
    run_end: int,
    checkpoints: np.ndarray,
    lag: int = 20,
) -> Dict[str, np.ndarray]:
    """Calculates the distance between the center atom and the neighbor atom
    in the checkpoint +/- lag time range.

    Args:
        nvt_run: An MDAnalysis ``Universe`` containing wrapped trajectory.
        center_atom: The center atom object.
        distance_dict: A dictionary of Cutoff distance of neighbor for each species.
        select_dict: A dictionary of atom species selection, where each atom species name is a key
            and the corresponding values are the selection language.
        run_start: Start frame of analysis.
        run_end: End frame of analysis.
        checkpoints: The frame numberings of interest to check for contiguous steps.
        lag: The range (+/- lag) of the contiguous steps. Default to 20.

    Returns:
        An array of distance between the center atom and the neighbor atoms
        in the checkpoint +/- lag time range.
    """
    coord_num: Dict[str, Union[List[List[int]], np.ndarray]] = {
        x: [[] for _ in range(lag * 2 + 1)] for x in distance_dict
    }
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    has = False
    for i, ts in enumerate(trj_analysis):
        log = False
        checkpoint = -1
        for j in checkpoints:
            if abs(i - j) <= lag:
                log = True
                has = True
                checkpoint = j
        if log:
            for kw in distance_dict:
                selection = select_shell(select_dict, distance_dict, center_atom, kw)
                shell = nvt_run.select_atoms(selection, periodic=True)
                coord_num[kw][i - checkpoint + lag].append(len(shell))
    one_atom_ave = {}
    if has:
        for kw in coord_num:
            np_arrays = np.array([np.array(time).mean() for time in coord_num[kw]])
            one_atom_ave[kw] = np_arrays
    return one_atom_ave

def heat_map(
    nvt_run: Universe,
    floating_atom: Atom,
    cluster_center_sites: List[int],
    cluster_terminal: Union[str, List[str]],
    cartesian_by_ref: np.ndarray,
    run_start: int,
    run_end: int,
    dim: str = "xyz",
) -> np.ndarray:
    """
    Calculates the heat map of the floating atom around the cluster. The coordinates are normalized to
    a cartesian coordinate system where the cluster_center_sites atom is the origin.

    Args:
        nvt_run: An MDAnalysis ``Universe`` containing wrapped trajectory.
        floating_atom: Floating atom species.
        cluster_center_sites: A list of nearest cluster center sites (atom index).
        cluster_terminal: The selection string for terminal atom species of the cluster
            (typically the binding site for the floating ion). The argument can be a str if
            all the terminal atoms have the same selection string and are equivalent, or a list
            if the terminal atoms are distinct and have different selection strings.
        cartesian_by_ref: Transformation matrix between cartesian and reference coordinate systems.
        run_start: Start frame of analysis.
        run_end: End frame of analysis.
        dim: Desired dimensions to calculate heat map. TODO: 2d support or dimension selection.

    Returns:
        The coordinates of the floating ion around clusters normalized to the desired cartesian coordinate system.
    """
    dimension = len(dim)
    if isinstance(cluster_terminal, list):
        mode = "ordered"
        if len(cluster_terminal) != dimension:
            term_order = ", ".join(c for c in dim)
            raise ValueError(f"Please specify the cluster_terminal in the order of {term_order}.")
    else:
        mode = "unordered"
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    coordinates = []
    for i, ts in enumerate(trj_analysis):
        if cluster_center_sites[i] == 0:
            pass
        else:
            center_atom = nvt_run.select_atoms("index " + str(cluster_center_sites[i]))[0]
            if mode == "ordered":
                selections = [
                    "(" + species + ") and (same resid as index " + str(center_atom.index) + ")"
                    for species in cluster_terminal
                ]
                bind_atoms_xyz = [nvt_run.select_atoms(sel, periodic=True) for sel in selections]
                vertex_atoms: List[Atom] = []
                for atoms in bind_atoms_xyz:
                    if len(atoms) == 1:
                        vertex_atoms.append(atoms[0])
                    elif len(atoms) > 1:
                        distances = distance_array(ts[floating_atom.index], atoms.positions, ts.dimensions)
                        idx = np.argpartition(distances[0], 1)
                        vertex_atoms.append(atoms[idx[0]])
                    else:
                        raise ValueError(
                            f"There should be at least 1 cluster_terminal atom in the {str(dim[i])} dimension."
                            f"Try broadening the selection at index {str(i + 1)} of the cluster_terminal "
                        )
            else:
                assert isinstance(cluster_terminal, str)
                selection = "(" + cluster_terminal + ") and (same resid as index " + str(center_atom.index) + ")"
                bind_atoms = nvt_run.select_atoms(selection, periodic=True)
                if len(bind_atoms) == dimension:
                    vertex_atoms = bind_atoms
                elif len(bind_atoms) > dimension:
                    distances = distance_array(ts[floating_atom.index], bind_atoms.positions, ts.dimensions)
                    idx = np.argpartition(distances[0], 3)
                    vertex_atoms = bind_atoms[idx[:3]]
                else:
                    raise ValueError(
                        f"There should be at least {dimension} cluster_terminal atoms in order to position "
                        "the floating ion. Try broadening the cluster_terminal selection"
                    )
            vector_a = atom_vec(vertex_atoms[0], center_atom, ts.dimensions)
            vector_b = atom_vec(vertex_atoms[1], center_atom, ts.dimensions)
            vector_c = atom_vec(vertex_atoms[2], center_atom, ts.dimensions)
            vector_atom = atom_vec(floating_atom, center_atom, ts.dimensions)
            basis_abc = np.transpose([vector_a, vector_b, vector_c])
            abc_atom = np.linalg.solve(basis_abc, vector_atom)
            unit_x = np.linalg.norm(
                cartesian_by_ref[0, 0] * vector_a
                + cartesian_by_ref[0, 1] * vector_b
                + cartesian_by_ref[0, 2] * vector_c
            )
            unit_y = np.linalg.norm(
                cartesian_by_ref[1, 0] * vector_a
                + cartesian_by_ref[1, 1] * vector_b
                + cartesian_by_ref[1, 2] * vector_c
            )
            unit_z = np.linalg.norm(
                cartesian_by_ref[2, 0] * vector_a
                + cartesian_by_ref[2, 1] * vector_b
                + cartesian_by_ref[2, 2] * vector_c
            )
            vector_x = cartesian_by_ref[0] / unit_x
            vector_y = cartesian_by_ref[1] / unit_y
            vector_z = cartesian_by_ref[2] / unit_z
            basis_xyz = np.transpose([vector_x, vector_y, vector_z])
            xyz_atom = np.linalg.solve(basis_xyz, abc_atom)
            coordinates.append(xyz_atom)
    return np.array(coordinates)

def process_evol(
    nvt_run: Universe,
    select_dict: Dict[str, str],
    in_list: Dict[str, List[np.ndarray]],
    out_list: Dict[str, List[np.ndarray]],
    distance_dict: Dict[str, float],
    run_start: int,
    run_end: int,
    lag: int,
    binding_cutoff: float,
    hopping_cutoff: float,
    smooth: int,
    cool: int,
    binding_site: str,
    center_atom: str,
):
    """Calculates the coordination number evolution of species around ``center_atom`` as a function of time,
    the coordination numbers are averaged over all frames around events when the center_atom
    hopping to and hopping out from the ``binding_site``.

    Args:
        nvt_run: An MDAnalysis ``Universe`` containing wrapped trajectory.
        select_dict: A dictionary of atom species selection, where each atom species name is a key
            and the corresponding values are the selection language.
        in_list: A list to store the distances for hopping in events.
        out_list: A list to store the distances for hopping out events.
        distance_dict: A dict of coordination cutoff distance of the neighbor species.
        run_start: Start frame of analysis.
        run_end: End frame of analysis.
        lag: The frame range (+/- lag) for check evolution.
        binding_cutoff: Binding cutoff distance.
        hopping_cutoff: Hopping out cutoff distance.
        smooth: The length of the smooth filter window. Default to 51.
        cool: The cool down frames between binding and hopping out.
        binding_site: The binding site of binding and hopping out events.
        center_atom: The solvation shell center atom.
    """
    center_atoms = nvt_run.select_atoms(select_dict.get(center_atom))
    for atom in tqdm(center_atoms[::]):
        neighbor_trj = neighbor_distance(
            nvt_run, atom, run_start + lag, run_end - lag, binding_site, select_dict, binding_cutoff
        )
        hopping_in, hopping_out = find_in_n_out(neighbor_trj, binding_cutoff, hopping_cutoff, smooth=smooth, cool=cool)
        if len(hopping_in) > 0:
            in_one = check_contiguous_steps(
                nvt_run,
                atom,
                distance_dict,
                select_dict,
                run_start,
                run_end,
                np.array(hopping_in) + lag,
                lag=lag,
            )
            for kw, value in in_one.items():
                in_list[kw].append(value)
        if len(hopping_out) > 0:
            out_one = check_contiguous_steps(
                nvt_run,
                atom,
                distance_dict,
                select_dict,
                run_start,
                run_end,
                np.array(hopping_out) + lag,
                lag=lag,
            )
            for kw, value in out_one.items():
                out_list[kw].append(value)

def get_full_coords(
    coords: np.ndarray,
    reflection: Optional[List[np.ndarray]] = None,
    rotation: Optional[List[np.ndarray]] = None,
    inversion: Optional[List[np.ndarray]] = None,
    sample: Optional[int] = None,
    dim: str = "xyz",
) -> np.ndarray:
    """
    A helper function for calculating the heatmap. It applies the ``reflection``, ``rotation`` and ``inversion``
    symmetry operations to ``coords`` and take ``sample`` number of samples.

    Args:
        coords: An array of coordinates.
        reflection: A list of reflection symmetry operation matrix.
        rotation: A list of rotation symmetry operation matrix.
        inversion: A list of inversion symmetry operation matrix.
        sample: Number of samples to take from ``coords``.
        dim: The dimensions of the coordinates. Default to "xyz".

    Returns:
        An array with ``sample`` number of coordinates.
    """
    dimension = len(dim)
    coords_full = coords
    if reflection:
        coords_copy = coords_full
        for mat in reflection:
            if mat.shape == (3,) and dimension == 3 or mat.shape == (2,) and dimension == 2:
                coords_ref = coords_copy * mat
                coords_full = np.concatenate((coords_full, coords_ref), axis=0)
            elif mat.shape == (3, 3) and dimension == 3 or mat.shape == (2, 2) and dimension == 2:
                coords_ref = np.dot(coords_copy, mat)
                coords_full = np.concatenate((coords_full, coords_ref), axis=0)
            else:
                raise ValueError(
                    f"Invalid reflection matrix. For {dimension}-D system, the matrix should be"
                    f" {dimension}x{dimension} or a vector of length {dimension}"
                )
    if rotation:
        coords_copy = coords_full
        for mat in rotation:
            if mat.shape == (3, 3) and dimension == 3 or mat.shape == (2, 2) and dimension == 2:
                coords_rot = np.dot(coords_copy, mat)
                coords_full = np.concatenate((coords_full, coords_rot), axis=0)
            else:
                raise ValueError(
                    f"Invalid rotation matrix. For {dimension}-D system, the matrix should be {dimension}x{dimension}."
                )
    if inversion:
        coords_copy = coords_full
        for mat in inversion:
            if mat.shape == (3, 3) and dimension == 3 or mat.shape == (2, 2) and dimension == 2:
                coords_inv = np.dot(coords_copy, mat)
                coords_full = np.concatenate((coords_full, coords_inv), axis=0)
            else:
                raise ValueError(
                    f"Invalid inversion matrix. For {dimension}-D system, the matrix should be {dimension}x{dimension}."
                )
    if sample:
        if coords_full.shape[0] > sample:
            index = np.random.choice(coords_full.shape[0], sample, replace=False)
            coords_full = coords_full[index]
        else:
            print(f"Warning: the number of coordinates < {sample}, will not perform sampling.")
    return coords_full

def cluster_coordinates(  # TODO: rewrite the method
    nvt_run: Universe,
    select_dict: Dict[str, str],
    run_start: int,
    run_end: int,
    species: List[str],
    distance: float,
    basis_vectors: Optional[Union[List[np.ndarray], np.ndarray]] = None,
    cluster_center: str = "center",
) -> np.ndarray:
    """Calculates the average position of a cluster.

    Args:
        nvt_run: An MDAnalysis ``Universe`` containing wrapped trajectory.
        select_dict: A dictionary of atom species selection, where each atom species name is a key
            and the corresponding values are the selection language.
        run_start: Start frame of analysis.
        run_end: End frame of analysis.
        species: A list of species in the cluster.
        distance: The coordination cutoff distance.
        basis_vectors: The basis vector for normalizing the coordinates of the cluster atoms.
        cluster_center: Cluster center atom species.

    Returns:
        An array of coordinates of the cluster atoms.
    """
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    cluster_center_atom = nvt_run.select_atoms(select_dict.get(cluster_center), periodic=True)[0]
    selection = (
        "("
        + " or ".join(s for s in species)
        + ") and (around "
        + str(distance)
        + " index "
        + str(cluster_center_atom.index)
        + ")"
    )
    shell = nvt_run.select_atoms(selection, periodic=True)
    cluster = []
    for atom in shell:
        coord_list = []
        for ts in trj_analysis:
            coord_list.append(atom.position)
        cluster.append(np.mean(np.array(coord_list), axis=0))
    cluster_array = np.array(cluster)
    if basis_vectors:
        if len(basis_vectors) == 2:
            vec1 = basis_vectors[0]
            vec2 = basis_vectors[1]
            vec3 = np.cross(vec1, vec2)
            vec2 = np.cross(vec1, vec3)
        elif len(basis_vectors) == 3:
            vec1 = basis_vectors[0]
            vec2 = basis_vectors[1]
            vec3 = basis_vectors[2]
        else:
            raise ValueError("incorrect vector format")
        vec1 = vec1 / np.linalg.norm(vec1)
        vec2 = vec2 / np.linalg.norm(vec2)
        vec3 = vec3 / np.linalg.norm(vec3)
        basis_xyz = np.transpose([vec1, vec2, vec3])
        cluster_norm = np.linalg.solve(basis_xyz, cluster_array.T).T
        cluster_norm = cluster_norm - np.mean(cluster_norm, axis=0)
        return cluster_norm
    return cluster_array

def num_of_neighbor(
    nvt_run: Universe,
    center_atom: Atom,
    distance_dict: Dict[str, float],
    select_dict: Dict[str, str],
    run_start,
    run_end,
    write=False,
    structure_code=None,
    write_freq=0,
    write_path=None,
) -> Dict[str, np.ndarray]:
    """Calculates the coordination number of each specified neighbor species and the total coordination number
    in the specified frame range.

    Args:
        nvt_run: An MDAnalysis ``Universe`` containing wrapped trajectory.
        center_atom: The solvation shell center atom.
        distance_dict: A dict of coordination cutoff distance of the neighbor species.
        select_dict: A dictionary of atom species selection, where each atom species name is a key
            and the corresponding values are the selection language.
        run_start: Start frame of analysis.
        run_end: End frame of analysis.
        write: Whether to writes out a series of desired solvation structures as ``*.xyz`` files.
        structure_code: An integer code representing the solvation structure to write out.
            For example, 221 is two species A, two species B and one species C.
        write_freq: Probability to write out files.
        write_path: Path to write out files.

    Returns:
        A diction containing the coordination number sequence of each specified neighbor species
        and the total coordination number sequence in the specified frame range .
    """
    time_count = 0
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    cn_values = {}
    species = list(distance_dict.keys())
    for kw in species:
        cn_values[kw] = np.zeros(int(len(trj_analysis)))
    cn_values["total"] = np.zeros(int(len(trj_analysis)))
    for ts in trj_analysis:
        digit_of_species = len(species) - 1
        for kw in species:
            selection = select_shell(select_dict, distance_dict, center_atom, kw)
            shell = nvt_run.select_atoms(selection, periodic=True)
            # for each atom in shell, create/add to dictionary
            # (key = atom id, value = list of values for step function)
            for _ in shell.atoms:
                cn_values[kw][time_count] += 1
                cn_values["total"][time_count] += 10 ** digit_of_species
            digit_of_species = digit_of_species - 1
        if write and cn_values["total"][time_count] == structure_code:
            a = np.random.random()
            if a > 1 - write_freq:
                print("writing")
                selection_write = " or ".join(
                    "(same resid as " + select_shell(select_dict, distance_dict, center_atom, kw) + ")"
                    for kw in species
                )
                center_selection = "same type as index " + str(center_atom.index)
                selection_write = "((" + selection_write + ") and not " + center_selection + ")"
                structure = nvt_run.select_atoms(selection_write, periodic=True)
                center_pos = ts[center_atom.index]
                center_name = center_atom.name
                path = write_path + str(center_atom.id) + "_" + str(int(ts.time)) + "_" + str(structure_code) + ".xyz"
                write_out(center_pos, center_name, structure, path)
        time_count += 1
    return cn_values

def num_of_neighbor_simple(
    nvt_run: Universe,
    center_atom: Atom,
    distance_dict: Dict[str, float],
    select_dict: Dict[str, str],
    run_start: int,
    run_end: int,
) -> Dict[str, np.ndarray]:
    """Calculates solvation structure type (1 for SSIP, 2 for CIP and 3 for AGG) with respect to the ``enter_atom``
    in the specified frame range.

    Args:
        nvt_run: An MDAnalysis ``Universe`` containing wrapped trajectory.
        center_atom: The solvation shell center atom.
        distance_dict: A dict of coordination cutoff distance of the neighbor species.
        select_dict: A dictionary of atom species selection, where each atom species name is a key
            and the corresponding values are the selection language.
        run_start: Start frame of analysis.
        run_end: End frame of analysis.

    Returns:
        A dict with "total" as the key and an array of the solvation structure type in the specified frame range
        as the value.
    """

    time_count = 0
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    center_selection = "same type as index " + str(center_atom.index)
    assert len(distance_dict) == 1, "Please only specify the counter-ion species in the distance_dict"
    species = list(distance_dict.keys())[0]
    cn_values = np.zeros(int(len(trj_analysis)))
    for ts in trj_analysis:
        selection = select_shell(select_dict, distance_dict, center_atom, species)
        shell = nvt_run.select_atoms(selection, periodic=True)
        shell_len = len(shell)
        if shell_len == 0:
            cn_values[time_count] = 1
        elif shell_len == 1:
            selection_species = select_shell(center_selection, distance_dict, shell.atoms[0], species)
            shell_species = nvt_run.select_atoms(selection_species, periodic=True)
            shell_species_len = len(shell_species) - 1
            if shell_species_len == 0:
                cn_values[time_count] = 2
            else:
                cn_values[time_count] = 3
        else:
            cn_values[time_count] = 3
        time_count += 1
    cn_values = {"total": cn_values}
    return cn_values

def angular_dist_of_neighbor(
    nvt_run: Universe,
    center_atom: Atom,
    distance_dict: Dict[str, float],
    select_dict: Dict[str, str],
    run_start: int,
    run_end: int,
    cip: bool = True,
) -> Dict[str, np.ndarray]:
    """
    Calculates the angle of a-c-b of center atom c in the specified frames.

    Args:
        nvt_run: An MDAnalysis ``Universe`` containing wrapped trajectory.
        center_atom: The center atom object.
        select_dict: A dictionary of atom species selection, where each atom species name is a key
            and the corresponding values are the selection language.
        distance_dict: A dict of coordination cutoff distance of the neighbor species. The key must be
            in the order of a, b, c, where a is the neighbor species used for determining coordination type,
            b is the other neighbor species, and c is the center species.
        run_start: Start frame of analysis.
        run_end: End frame of analysis.
        cip: Only includes contact ion pair structures with only one `a` and one `c` atoms.
            Default to True.

    Returns:
        An array of angles of a-c-b occurrence in the specified frames.
    """
    names = list(distance_dict.keys())
    assert len(names) == 3, "Invalid number of keys in distance_dict, should be 3."
    neighbor_a, neighbor_b, center_c = tuple(names)
    acb_angle = []
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    for ts in trj_analysis:
        a_selection = select_shell(select_dict, distance_dict, center_atom, neighbor_a)
        a_group = nvt_run.select_atoms(a_selection, periodic=True)
        a_num = len(a_group)
        if a_num == 0:
            continue
        if a_num == 1:
            c_selection = select_shell(select_dict, distance_dict, a_group.atoms[0], center_c)
            c_atoms = nvt_run.select_atoms(c_selection, periodic=True)
            shell_species_len = len(c_atoms) - 1
            if shell_species_len == 0:
                shell_type = "cip"
            else:
                shell_type = "agg"
        else:
            shell_type = "agg"
        if shell_type == "agg" and cip:
            continue
        c_pos = center_atom.position
        for a_atom in a_group.atoms:
            a_pos = a_atom.position
            b_selection = select_shell(select_dict, distance_dict, center_atom, neighbor_b)
            b_group = nvt_run.select_atoms(b_selection, periodic=True)
            for b_atom in b_group.atoms:
                b_pos = b_atom.position
                theta = angle(a_pos, c_pos, b_pos)
                acb_angle.append(theta)
    return {"total": np.array(acb_angle)}

def num_of_neighbor_specific(
    nvt_run: Universe,
    center_atom: Atom,
    distance_dict: Dict[str, float],
    select_dict: Dict[str, str],
    run_start: int,
    run_end: int,
    counter_atom: str = "anion",
) -> Dict[str, np.ndarray]:
    """
    Calculates the coordination number of each specific solvation structure type (SSIP, CIP, AGG).

    Args:
        nvt_run: An MDAnalysis ``Universe`` containing wrapped trajectory.
        center_atom: The center atom object.
        distance_dict: A dict of coordination cutoff distance of the neighbor species.
        select_dict: A dictionary of atom species selection, where each atom species name is a key
            and the corresponding values are the selection language.
        run_start: Start frame of analysis.
        run_end: End frame of analysis.
        counter_atom: The neighbor counter ion species. Default to "anion".

    Returns:
        A tuple containing three dictionary of the coordination number of each neighbor species
        and total coordination number for the three solvation structure type, respectively.
    """
    time_count = 0
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    cip_step = []
    ssip_step = []
    agg_step = []
    cn_values = {}
    for kw in distance_dict:
        cn_values[kw] = np.zeros(int(len(trj_analysis)))
    cn_values["total"] = np.zeros(int(len(trj_analysis)))
    for ts in trj_analysis:
        for kw in distance_dict:
            kw_selection = select_shell(select_dict, distance_dict, center_atom, kw)
            kw_shell = nvt_run.select_atoms(kw_selection, periodic=True)
            cn_values[kw][time_count] += len(kw_shell)
            cn_values["total"][time_count] += len(kw_shell)

        selection = select_shell(select_dict, distance_dict, center_atom, counter_atom)
        shell = nvt_run.select_atoms(selection, periodic=True)
        shell_len = len(shell)
        center_selection = "same type as index " + str(center_atom.index)
        if shell_len == 0:
            ssip_step.append(time_count)
        elif shell_len == 1:
            selection_species = select_shell(center_selection, distance_dict, shell.atoms[0], counter_atom)
            shell_species = nvt_run.select_atoms(selection_species, periodic=True)
            shell_species_len = len(shell_species) - 1
            if shell_species_len == 0:
                cip_step.append(time_count)
            else:
                agg_step.append(time_count)
        else:
            agg_step.append(time_count)
        time_count += 1
    cn_dict = {}
    for kw in distance_dict:
        cn_dict["ssip_" + kw] = cn_values[kw][ssip_step]
        cn_dict["cip_" + kw] = cn_values[kw][cip_step]
        cn_dict["agg_" + kw] = cn_values[kw][agg_step]
    return cn_dict

def full_solvation_structure(  # TODO: rewrite the method
    nvt_run: Universe,
    center_atom: Atom,
    center_species: str,
    counter_species: str,
    select_dict: Dict[str, str],
    distance: float,
    run_start: int,
    run_end: int,
    depth: int = 4,
) -> np.ndarray:
    """
    Obtain the solvation structure of a full connected ion network with depth-first traversal.

    Args:
        nvt_run: An MDAnalysis ``Universe`` containing wrapped trajectory.
        center_atom: The center atom object.
        center_species: The center ion species. It should be the atom directly connect to the counter ion.
        counter_species: The neighbor counter ion species. It should be the atom directly connect to the center ion.
        select_dict: A dictionary of atom species selection, where each atom species name is a key
            and the corresponding values are the selection language.
        distance: The coordination cutoff distance.
        run_start: Start frame of analysis.
        run_end: End frame of analysis.
        depth: The depth of the traversal. Default to 4

    Returns:
        Return an array of full solvation structure of the specified frames. Each solvation structure is represented by
        an array of the number of ions from the first to the n-th solvation shell with n=``depth``.
    """
    center_selection = select_dict.get(center_species)
    counter_selection = select_dict.get(counter_species)
    assert (center_selection is not None) and (counter_selection is not None)

    def select_counter_ion(selection, dist, atom):
        return "(" + selection + " and around " + str(dist) + " same fragment as index " + str(atom.index) + ")"

    def center_shell(this_shell, this_layer, frame):
        for counter in this_shell.atoms:
            if counter.id not in counter_ion_list:
                counter_ion_list.append(counter.id)
                cn_values[frame][this_layer] += 1
                if this_layer + 1 < depth:
                    next_shell = nvt_run.select_atoms(
                        select_counter_ion(center_selection, distance, counter),
                        periodic=True,
                    )
                    counter_shell(next_shell, this_layer + 1, frame)

    def counter_shell(this_shell, this_layer, frame):
        for center in this_shell.atoms:
            if center.id not in center_ion_list:
                center_ion_list.append(center.id)
                cn_values[frame][this_layer] += 1
                if this_layer + 1 < depth:
                    next_shell = nvt_run.select_atoms(
                        select_counter_ion(counter_selection, distance, center),
                        periodic=True,
                    )
                    center_shell(next_shell, this_layer + 1, frame)

    time_count = 0
    trj_analysis = nvt_run.trajectory[run_start:run_end:]
    cn_values = np.zeros((int(len(trj_analysis)), depth))
    for ts in trj_analysis:
        center_ion_list: List[np.int_] = [center_atom.id]
        counter_ion_list: List[np.int_] = []
        first_shell = nvt_run.select_atoms(
            select_counter_ion(counter_selection, distance, center_atom),
            periodic=True,
        )
        center_shell(first_shell, 0, time_count)
    return cn_values

def concat_coord_array(
    nvt_run: Universe,
    func: Callable,
    center_atoms: AtomGroup,
    distance_dict: Dict[str, float],
    select_dict: Dict[str, str],
    run_start: int,
    run_end: int,
    **kwargs: Union[bool, str],
) -> Dict[str, np.ndarray]:
    """
    A helper function to analyze the coordination number/structure of every atoms in an ``AtomGroup`` using the
    specified function.

    Args:
        nvt_run: An MDAnalysis ``Universe`` containing wrapped trajectory.
        func: One of the neighbor statistical method (num_of_neighbor, num_of_neighbor_simple)
        center_atoms: Atom group of the center atoms.
        distance_dict: A dictionary of coordination cutoff distance of the neighbor species.
        select_dict: A dictionary of atom species selection, where each atom species name is a key
            and the corresponding values are the selection language.
        run_start: Start frame of analysis.
        run_end: End frame of analysis.

    Returns:
        A diction containing the coordination number sequence of each specified neighbor species
        and the total coordination number sequence in the specified frame range.
    """
    num_array = func(nvt_run, center_atoms[0], distance_dict, select_dict, run_start, run_end, **kwargs)
    for atom in tqdm(center_atoms[1::]):
        this_atom = func(nvt_run, atom, distance_dict, select_dict, run_start, run_end, **kwargs)
        for kw in num_array:
            num_array[kw] = np.concatenate((num_array.get(kw), this_atom.get(kw)), axis=0)
    return num_array

def write_out(center_pos: np.ndarray, center_name: str, neighbors: AtomGroup, path: str):
    """
    Helper function for solvation structure coordinates write out.

    Args:
        center_pos: The coordinates of the center atom in the frame.
        center_name: The element name of the center atom in the frame.
        neighbors: The neighbor AtomGroup.
        path: The path to write out ``*.xyz`` file.
    """
    lines = []
    lines.append(str(len(neighbors) + 1))
    lines.append("")
    lines.append(f"{center_name} 0.0000000 0.0000000 0.0000000")
    box = neighbors.dimensions
    half_box = np.array([box[0], box[1], box[2]]) / 2
    for atom in neighbors:
        locs = []
        for i in range(3):
            loc = atom.position[i] - center_pos[i]
            if loc > half_box[i]:
                loc = loc - box[i]
            elif loc < -half_box[i]:
                loc = loc + box[i]
            else:
                pass
            locs.append(loc)
        element_name = atom.name
        assert element_name is not None
        line = element_name + " " + " ".join(str(loc) for loc in locs)
        lines.append(line)
    with open(path, "w") as xyz_file:
        xyz_file.write("\n".join(lines))

def select_shell(
    select: Union[Dict[str, str], str], distance: Union[Dict[str, float], str], center_atom: Atom, kw: str
) -> str:
    """
    Select a group of atoms that is within a distance of an ``center_atom``.

    Args:
        select: A selection string of neighbors or a dictionary of atom species selection, where each atom
            species name is a key and the corresponding values are the selection string.
        distance: A neighbor cutoff distance or a dict of cutoff distances of neighbor species.
        center_atom: The solvation shell center ``Atom`` object
        kw: The key for the select and/or distance dictionary if applicable.

    Returns:
        A selection string specifying the neighbor species within a distance of the ``center_atom``.
    """
    if isinstance(select, dict):
        species_selection = select[kw]
        if species_selection is None:
            raise ValueError("Species specified does not match entries in the select dict.")
    else:
        species_selection = select
    if isinstance(distance, dict):
        distance_value = distance[kw]
        if distance_value is None:
            raise ValueError("Species specified does not match entries in the distance dict.")
        distance_str = str(distance_value)
    else:
        distance_str = distance
    selection = "(" + species_selection + ") and (around " + distance_str + " index " + str(center_atom.index) + ")"
    return selection


import logging
import string
import sys
from collections import deque

import numpy as np
from scipy.special import logsumexp
from sklearn.base import BaseEstimator
from sklearn.utils import check_array, check_random_state
import scipy.optimize
from . import _hmmc, _utils
from .utils import normalize, log_normalize, iter_from_X_lengths, log_mask_zero

_log = logging.getLogger(__name__)
#: Supported decoder algorithms.
DECODER_ALGORITHMS = frozenset(("viterbi", "map"))

class ConvergenceMonitor:
    """Monitors and reports convergence to :data:`sys.stderr`.

    Parameters
    ----------
    tol : double
        Convergence threshold. EM has converged either if the maximum
        number of iterations is reached or the log probability
        improvement between the two consecutive iterations is less
        than threshold.

    n_iter : int
        Maximum number of iterations to perform.

    verbose : bool
        If ``True`` then per-iteration convergence reports are printed,
        otherwise the monitor is mute.

    Attributes
    ----------
    history : deque
        The log probability of the data for the last two training
        iterations. If the values are not strictly increasing, the
        model did not converge.

    iter : int
        Number of iterations performed while training the model.

    Examples
    --------
    Use custom convergence criteria by subclassing ``ConvergenceMonitor``
    and redefining the ``converged`` method. The resulting subclass can
    be used by creating an instance and pointing a model's ``monitor_``
    attribute to it prior to fitting.

    >>> from hmmlearn.base import ConvergenceMonitor
    >>> from hmmlearn import hmm
    >>>
    >>> class ThresholdMonitor(ConvergenceMonitor):
    ...     @property
    ...     def converged(self):
    ...         return (self.iter == self.n_iter or
    ...                 self.history[-1] >= self.tol)
    >>>
    >>> model = hmm.GaussianHMM(n_components=2, tol=5, verbose=True)
    >>> model.monitor_ = ThresholdMonitor(model.monitor_.tol,
    ...                                   model.monitor_.n_iter,
    ...                                   model.monitor_.verbose)
    """
    _template = "{iter:>10d} {logprob:>16.4f} {delta:>+16.4f}"

    def __init__(self, tol, n_iter, verbose):
        self.tol = tol
        self.n_iter = n_iter
        self.verbose = verbose
        self.history = deque(maxlen=2)
        self.iter = 0

    def __repr__(self):
        class_name = self.__class__.__name__
        params = sorted(dict(vars(self), history=list(self.history)).items())
        return ("{}(\n".format(class_name)
                + "".join(map("    {}={},\n".format, *zip(*params)))
                + ")")

    def _reset(self):
        """Reset the monitor's state."""
        self.iter = 0
        self.history.clear()

    def report(self, logprob):
        """Reports convergence to :data:`sys.stderr`.

        The output consists of three columns: iteration number, log
        probability of the data at the current iteration and convergence
        rate.  At the first iteration convergence rate is unknown and
        is thus denoted by NaN.

        Parameters
        ----------
        logprob : float
            The log probability of the data as computed by EM algorithm
            in the current iteration.
        """
        if self.verbose:
            delta = logprob - self.history[-1] if self.history else np.nan
            message = self._template.format(
                iter=self.iter + 1, logprob=logprob, delta=delta)
            print(message, file=sys.stderr)

        self.history.append(logprob)
        self.iter += 1

    @property
    def converged(self):
        """``True`` if the EM algorithm converged and ``False`` otherwise."""
        # XXX we might want to check that ``logprob`` is non-decreasing.
        return (self.iter == self.n_iter or
                (len(self.history) == 2 and
                 self.history[1] - self.history[0] < self.tol))

class _BaseHMM(BaseEstimator):
    r"""Base class for Hidden Markov Models.

    This class allows for easy evaluation of, sampling from, and
    maximum a posteriori estimation of the parameters of a HMM.

    See the instance documentation for details specific to a
    particular object.

    Parameters
    ----------
    n_components : int
        Number of states in the model.

    startprob_prior : array, shape (n_components, ), optional
        Parameters of the Dirichlet prior distribution for
        :attr:`startprob_`.

    transmat_prior : array, shape (n_components, n_components), optional
        Parameters of the Dirichlet prior distribution for each row
        of the transition probabilities :attr:`transmat_`.

    algorithm : string, optional
        Decoder algorithm. Must be one of "viterbi" or "map".
        Defaults to "viterbi".

    random_state: RandomState or an int seed, optional
        A random number generator instance.

    n_iter : int, optional
        Maximum number of iterations to perform.

    tol : float, optional
        Convergence threshold. EM will stop if the gain in log-likelihood
        is below this value.

    verbose : bool, optional
        When ``True`` per-iteration convergence reports are printed
        to :data:`sys.stderr`. You can diagnose convergence via the
        :attr:`monitor_` attribute.

    params : string, optional
        Controls which parameters are updated in the training
        process.  Can contain any combination of 's' for startprob,
        't' for transmat, and other characters for subclass-specific
        emission parameters. Defaults to all parameters.

    init_params : string, optional
        Controls which parameters are initialized prior to
        training.  Can contain any combination of 's' for
        startprob, 't' for transmat, and other characters for
        subclass-specific emission parameters. Defaults to all
        parameters.

    Attributes
    ----------
    monitor\_ : ConvergenceMonitor
        Monitor object used to check the convergence of EM.

    startprob\_ : array, shape (n_components, )
        Initial state occupation distribution.

    transmat\_ : array, shape (n_components, n_components)
        Matrix of transition probabilities between states.
    """
    def __init__(self, n_components=1,
                 startprob_prior=1.0, transmat_prior=1.0,
                 algorithm="viterbi", random_state=None,
                 n_iter=10, tol=1e-2, verbose=False,
                 params=string.ascii_letters,
                 init_params=string.ascii_letters,
                 A1 = None,
                 a = 0,
                 b = 1,
                 grad_iter = 100,
                 grad_conv = 10**-2,
                 grad_lr = 10**-2,
                 t0 = 0,
                 transmat_ = None,
                 startprob_ = None):
        self.n_components = n_components
        self.params = params
        self.init_params = init_params
        self.startprob_prior = startprob_prior
        self.transmat_prior = transmat_prior
        self.algorithm = algorithm
        self.random_state = random_state
        self.n_iter = n_iter
        self.tol = tol
        self.verbose = verbose
        self.monitor_ = ConvergenceMonitor(self.tol, self.n_iter, self.verbose)
        # added samuelbray32
        self.A1 = A1 #time dependent transition matrix
        self.a = a #parameters of sigmoid function
        self.b = b
        self.grad_iter = grad_iter #stop criteria for _grad_descent
        self.grad_conv = grad_conv
        self.grad_lr = grad_lr
        self.t0 = t0
        self.transmat_ = transmat_
        self.startprob_ = startprob_
        self.track_params = []
        self.grad_method = 'newtons_linesearch'

    def get_stationary_distribution(self):
        """Compute the stationary distribution of states.
        """
        # The stationary distribution is proportional to the left-eigenvector
        # associated with the largest eigenvalue (i.e., 1) of the transition
        # matrix.
        _utils.check_is_fitted(self, "transmat_")
        eigvals, eigvecs = np.linalg.eig(self.transmat_.T)
        eigvec = np.real_if_close(eigvecs[:, np.argmax(eigvals)])
        return eigvec / eigvec.sum()

    def score_samples(self, X, lengths=None):
        """Compute the log probability under the model and compute posteriors.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.

        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.

        Returns
        -------
        logprob : float
            Log likelihood of ``X``.

        posteriors : array, shape (n_samples, n_components)
            State-membership probabilities for each sample in ``X``.

        See Also
        --------
        score : Compute the log probability under the model.
        decode : Find most likely state sequence corresponding to ``X``.
        """
        _utils.check_is_fitted(self, "startprob_")
        self._check()

        X = check_array(X)
        n_samples = X.shape[0]
        logprob = 0
        posteriors = np.zeros((n_samples, self.n_components))
        for i, j in iter_from_X_lengths(X, lengths):
            framelogprob = self._compute_log_likelihood(X[i:j])
            logprobij, fwdlattice = self._do_forward_pass(framelogprob)
            logprob += logprobij

            bwdlattice = self._do_backward_pass(framelogprob)
            posteriors[i:j] = self._compute_posteriors(fwdlattice, bwdlattice)
        return logprob, posteriors

    def score(self, X, lengths=None):
        """Compute the log probability under the model.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.

        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.

        Returns
        -------
        logprob : float
            Log likelihood of ``X``.

        See Also
        --------
        score_samples : Compute the log probability under the model and
            posteriors.
        decode : Find most likely state sequence corresponding to ``X``.
        """
        _utils.check_is_fitted(self, "startprob_")
        self._check()

        X = check_array(X)
        # XXX we can unroll forward pass for speed and memory efficiency.
        logprob = 0
        for i, j in iter_from_X_lengths(X, lengths):
            framelogprob = self._compute_log_likelihood(X[i:j])
            logprobij, _fwdlattice = self._do_forward_pass(framelogprob)
            logprob += logprobij
        return logprob

    def _decode_viterbi(self, X):
        framelogprob = self._compute_log_likelihood(X)
        return self._do_viterbi_pass(framelogprob)

    def _decode_map(self, X):
        _, posteriors = self.score_samples(X)
        logprob = np.max(posteriors, axis=1).sum()
        state_sequence = np.argmax(posteriors, axis=1)
        return logprob, state_sequence

    def decode(self, X, lengths=None, algorithm=None):
        """Find most likely state sequence corresponding to ``X``.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.

        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.

        algorithm : string
            Decoder algorithm. Must be one of "viterbi" or "map".
            If not given, :attr:`decoder` is used.

        Returns
        -------
        logprob : float
            Log probability of the produced state sequence.

        state_sequence : array, shape (n_samples, )
            Labels for each sample from ``X`` obtained via a given
            decoder ``algorithm``.

        See Also
        --------
        score_samples : Compute the log probability under the model and
            posteriors.
        score : Compute the log probability under the model.
        """
        _utils.check_is_fitted(self, "startprob_")
        self._check()

        algorithm = algorithm or self.algorithm
        if algorithm not in DECODER_ALGORITHMS:
            raise ValueError("Unknown decoder {!r}".format(algorithm))

        decoder = {
            "viterbi": self._decode_viterbi,
            "map": self._decode_map
        }[algorithm]

        X = check_array(X)
        n_samples = X.shape[0]
        logprob = 0
        state_sequence = np.empty(n_samples, dtype=int)
        for i, j in iter_from_X_lengths(X, lengths):
            # XXX decoder works on a single sample at a time!
            logprobij, state_sequenceij = decoder(X[i:j])
            logprob += logprobij
            state_sequence[i:j] = state_sequenceij

        return logprob, state_sequence

    def predict(self, X, lengths=None):
        """Find most likely state sequence corresponding to ``X``.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.

        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.

        Returns
        -------
        state_sequence : array, shape (n_samples, )
            Labels for each sample from ``X``.
        """
        _, state_sequence = self.decode(X, lengths)
        return state_sequence

    def predict_proba(self, X, lengths=None):
        """Compute the posterior probability for each state in the model.

        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.

        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.

        Returns
        -------
        posteriors : array, shape (n_samples, n_components)
            State-membership probabilities for each sample from ``X``.
        """
        _, posteriors = self.score_samples(X, lengths)
        return posteriors

    def sample(self, n_samples=1, random_state=None):
        """Generate random samples from the model.

        Parameters
        ----------
        n_samples : int
            Number of samples to generate.

        random_state : RandomState or an int seed
            A random number generator instance. If ``None``, the object's
            ``random_state`` is used.

        Returns
        -------
        X : array, shape (n_samples, n_features)
            Feature matrix.

        state_sequence : array, shape (n_samples, )
            State sequence produced by the model.
        """
        _utils.check_is_fitted(self, "startprob_")
        self._check()

        if random_state is None:
            random_state = self.random_state
        random_state = check_random_state(random_state)

        startprob_cdf = np.cumsum(self.startprob_)
        transmat_cdf = np.cumsum(self.transmat_, axis=1)

        currstate = (startprob_cdf > random_state.rand()).argmax()
        state_sequence = [currstate]
        X = [self._generate_sample_from_state(
            currstate, random_state=random_state)]

        for t in range(n_samples - 1):
            currstate = (transmat_cdf[currstate] > random_state.rand()) \
                .argmax()
            state_sequence.append(currstate)
            X.append(self._generate_sample_from_state(
                currstate, random_state=random_state))

        return np.atleast_2d(X), np.array(state_sequence, dtype=int)

    def fit(self, X, lengths=None):
        """Estimate model parameters.

        An initialization step is performed before entering the
        EM algorithm. If you want to avoid this step for a subset of
        the parameters, pass proper ``init_params`` keyword argument
        to estimator's constructor.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.

        lengths : array-like of integers, shape (n_sequences, )
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.

        Returns
        -------
        self : object
            Returns self.
        """
        X = check_array(X)
        self._init(X, lengths=lengths)
        self._check()

        self.monitor_._reset()
        for iter in range(self.n_iter):
            print('E iteration: ', iter)
            stats = self._initialize_sufficient_statistics()
            curr_logprob = 0
            for i, j in iter_from_X_lengths(X, lengths):
                framelogprob = self._compute_log_likelihood(X[i:j])
                logprob, fwdlattice = self._do_forward_pass(framelogprob)
                curr_logprob += logprob
                bwdlattice = self._do_backward_pass(framelogprob)
                posteriors = self._compute_posteriors(fwdlattice, bwdlattice)
                self._accumulate_sufficient_statistics(
                    stats, X[i:j], framelogprob, posteriors, fwdlattice,
                    bwdlattice)
                print('fwd: ', fwdlattice[10])
                print('bwd: ', bwdlattice[10])
            # XXX must be before convergence check, because otherwise
            #     there won't be any updates for the case ``n_iter=1``.
            del fwdlattice
            del bwdlattice
            del framelogprob
            self._do_mstep(stats, posteriors)

            self.monitor_.report(curr_logprob)
            if self.monitor_.converged:
                break

        if (self.transmat_.sum(axis=1) == 0).any():
            _log.warning("Some rows of transmat_ have zero sum because no "
                         "transition from the state was ever observed.")

        return self

    def _do_viterbi_pass(self, framelogprob):
        n_samples, n_components = framelogprob.shape
        state_sequence, logprob = _hmmc._viterbi(
            n_samples, n_components, log_mask_zero(self.startprob_),
            log_mask_zero(self.transmat_), framelogprob)
        return logprob, state_sequence

    def _do_forward_pass(self, framelogprob):
        n_samples, n_components = framelogprob.shape
        fwdlattice = np.zeros((n_samples, n_components))
        # _hmmc._forward(n_samples, n_components,
        #                log_mask_zero(self.startprob_),
        #                self.transmat_,
        #                framelogprob, fwdlattice,
        #                self.A1,
        #                self.U(np.arange(n_samples)+self.t0))
        t = np.arange(n_samples)+self.t0
        log_trans = self.transmat_ + np.transpose(self.A1[:,:,None]*self.U(t),[2,0,1])
        log_trans[log_trans<0] = 10**-20
        print('trans: ', np.min(log_trans))
        log_trans = log_mask_zero(log_trans)
        print('log_trans: ', np.max(log_trans))
        _hmmc._forward(n_samples, n_components,
                       log_mask_zero(self.startprob_),
                       log_trans,
                       framelogprob, fwdlattice,)

        with np.errstate(under="ignore"):
            return logsumexp(fwdlattice[-1]), fwdlattice

    def _do_backward_pass(self, framelogprob):
        n_samples, n_components = framelogprob.shape
        bwdlattice = np.zeros((n_samples, n_components))
        # _hmmc._backward(n_samples, n_components,
        #                 log_mask_zero(self.startprob_),
        #                 self.transmat_,
        #                 framelogprob, bwdlattice,
        #                 self.A1,
        #                 self.U(np.arange(n_samples)+self.t0))
        t = np.arange(n_samples)+self.t0
        log_trans = self.transmat_ + np.transpose(self.A1[:,:,None]*self.U(t),[2,0,1])
        log_trans[log_trans<0] = 10**-20
        print('trans: ', np.min(log_trans))
        log_trans = log_mask_zero(log_trans)
        print('log_trans: ', np.max(log_trans))
        _hmmc._backward(n_samples, n_components,
                       log_mask_zero(self.startprob_),
                       log_trans,
                       framelogprob, bwdlattice,)
        return bwdlattice

    def _compute_posteriors(self, fwdlattice, bwdlattice):
        # gamma is guaranteed to be correctly normalized by logprob at
        # all frames, unless we do approximate inference using pruning.
        # So, we will normalize each frame explicitly in case we
        # pruned too aggressively.
        log_gamma = fwdlattice + bwdlattice
        log_normalize(log_gamma, axis=1)
        with np.errstate(under="ignore"):
            return np.exp(log_gamma)

    def _init(self, X, lengths):
        """Initializes model parameters prior to fitting.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.

        lengths : array-like of integers, shape (n_sequences, )
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``.
        """
        init = 1. / self.n_components
        if 's' in self.init_params or not hasattr(self, "startprob_"):
            self.startprob_ = np.full(self.n_components, init)
        if 't' in self.init_params or not hasattr(self, "transmat_"):
            self.transmat_ = np.full((self.n_components, self.n_components),
                                     init)
        n_fit_scalars_per_param = self._get_n_fit_scalars_per_param()
        n_fit_scalars = sum(n_fit_scalars_per_param[p] for p in self.params)
        if X.size < n_fit_scalars:
            _log.warning("Fitting a model with {} free scalar parameters with "
                         "only {} data points will result in a degenerate "
                         "solution.".format(n_fit_scalars, X.size))

    def _check(self):
        """Validates model parameters prior to fitting.

        Raises
        ------

        ValueError
            If any of the parameters are invalid, e.g. if :attr:`startprob_`
            don't sum to 1.
        """
        self.startprob_ = np.asarray(self.startprob_)
        if len(self.startprob_) != self.n_components:
            raise ValueError("startprob_ must have length n_components")
        if not np.allclose(self.startprob_.sum(), 1.0):
            raise ValueError("startprob_ must sum to 1.0 (got {:.4f})"
                             .format(self.startprob_.sum()))

        self.transmat_ = np.asarray(self.transmat_)
        if self.transmat_.shape != (self.n_components, self.n_components):
            raise ValueError(
                "transmat_ must have shape (n_components, n_components)")
        if not np.allclose(self.transmat_.sum(axis=1), 1.0):
            raise ValueError("rows of transmat_ must sum to 1.0 (got {})"
                             .format(self.transmat_.sum(axis=1)))

    def _compute_log_likelihood(self, X):
        """Computes per-component log probability under the model.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.

        Returns
        -------
        logprob : array, shape (n_samples, n_components)
            Log probability of each sample in ``X`` for each of the
            model states.
        """

    def _generate_sample_from_state(self, state, random_state=None):
        """Generates a random sample from a given component.

        Parameters
        ----------
        state : int
            Index of the component to condition on.

        random_state: RandomState or an int seed
            A random number generator instance. If ``None``, the object's
            ``random_state`` is used.

        Returns
        -------
        X : array, shape (n_features, )
            A random sample from the emission distribution corresponding
            to a given component.
        """

    # Methods used by self.fit()

    def _initialize_sufficient_statistics(self):
        """Initializes sufficient statistics required for M-step.

        The method is *pure*, meaning that it doesn't change the state of
        the instance.  For extensibility computed statistics are stored
        in a dictionary.

        Returns
        -------
        nobs : int
            Number of samples in the data.

        start : array, shape (n_components, )
            An array where the i-th element corresponds to the posterior
            probability of the first sample being generated by the i-th
            state.

        trans : array, shape (n_components, n_components)
            An array where the (i, j)-th element corresponds to the
            posterior probability of transitioning between the i-th to j-th
            states.
        """
        stats = {'nobs': 0,
                 'start': np.zeros(self.n_components),
                 'trans': np.zeros((self.n_components, self.n_components))}
        return stats

    def _accumulate_sufficient_statistics(self, stats, X, framelogprob,
                                          posteriors, fwdlattice, bwdlattice):
        """Updates sufficient statistics from a given sample.

        Parameters
        ----------
        stats : dict
            Sufficient statistics as returned by
            :meth:`~base._BaseHMM._initialize_sufficient_statistics`.

        X : array, shape (n_samples, n_features)
            Sample sequence.

        framelogprob : array, shape (n_samples, n_components)
            Log-probabilities of each sample under each of the model states.

        posteriors : array, shape (n_samples, n_components)
            Posterior probabilities of each sample being generated by each
            of the model states.

        fwdlattice, bwdlattice : array, shape (n_samples, n_components)
            Log-forward and log-backward probabilities.
        """
        stats['nobs'] += 1
        if 's' in self.params:
            stats['start'] += posteriors[0]
        if 't' in self.params:
            n_samples, n_components = framelogprob.shape
            # when the sample is of length 1, it contains no transitions
            # so there is no reason to update our trans. matrix estimate
            if n_samples <= 1:
                return

            log_xi_sum = np.full((n_components, n_components), -np.inf)
            _hmmc._compute_log_xi_sum(n_samples, n_components, fwdlattice,
                                      log_mask_zero(self.transmat_),
                                      bwdlattice, framelogprob,
                                      log_xi_sum)
            with np.errstate(under="ignore"):
                stats['trans'] += np.exp(log_xi_sum)

    def _do_mstep(self, stats, posteriors):
        """Performs the M-step of EM algorithm.

        Parameters
        ----------
        stats : dict
            Sufficient statistics updated from all available samples.
        """
        # If a prior is < 1, `prior - 1 + starts['start']` can be negative.  In
        # that case maximization of (n1+e1) log p1 + ... + (ns+es) log ps under
        # the conditions sum(p) = 1 and all(p >= 0) show that the negative
        # terms can just be set to zero.
        # The ``np.where`` calls guard against updating forbidden states
        # or transitions in e.g. a left-right HMM.
        if 's' in self.params:
            startprob_ = np.maximum(self.startprob_prior - 1 + stats['start'],
                                    0)
            self.startprob_ = np.where(self.startprob_ == 0, 0, startprob_)
            normalize(self.startprob_)
        if 't' in self.params:
            transmat_ = np.maximum(self.transmat_prior - 1 + stats['trans'], 0)
            self.transmat_ = np.where(self.transmat_ == 0, 0, transmat_)
            normalize(self.transmat_, axis=1)
        if not self.a == None:
            # self._grad_ascent(self._transition_posterior(posteriors))
            log_P_tij = log_mask_zero(np.zeros((posteriors.shape[0],self.n_components,self.n_components)))
            _hmmc._transition_posterior(log_P_tij.shape[0], self.n_components, log_P_tij, log_mask_zero(posteriors))
            if self.grad_method == 'newtons_linesearch':
                self._newtons_linesearch(np.exp(log_P_tij[1:]))
            elif self.grad_method == 'newtons':
                self._newtons(np.exp(log_P_tij[1:]))
            elif self.grad_method == 'grad_ascent':
                self._grad_ascent(np.exp(log_P_tij[1:]))
            elif self.grad_method == 'grad_ascent_linesearch':
                self._grad_ascent_linesearch(np.exp(log_P_tij[1:]))
            elif self.grad_method == 'scipy':
                self._scipy(np.exp(log_P_tij[1:]))
            else:
                print (self.grad_method, ' is not a defined method')

    def _transition_posterior_external(self, log_P_tij, posteriors):
        _hmmc._transition_posterior(log_P_tij.shape[0], self.n_components, log_P_tij, log_mask_zero(posteriors))

    def U(self, t):
        return 1-self.sigma(t)

    def dU(self, t):
        return -self.sigma(t)*(1-self.sigma(t))

    # NEW VERSION: x = t/a - b
    def sigma(self, t):
        # print('new sigma')
        t = t/24/60/120
        return (1+np.exp(-(t-self.b)/self.a))**-1

    def _grad_ascent(self, P_tij):
        print('transition posteriors:', P_tij.shape)
        print('max: ', np.max(P_tij))
        print('sumcheck: ', np.sum(P_tij[1,:,:]), np.sum(P_tij[-1,:,:]) )
        if np.array(self.grad_lr).size == 1:
            self.grad_lr = [self.grad_lr, self.grad_lr]
        t = np.arange(P_tij.shape[0]) + self.t0
        for iter in range(self.grad_iter):
            mu = self._mu(t) * P_tij
            db = -mu.sum()/self.a *self.grad_lr[1]
            da = -(((t/24/60/120-self.b)[:,None,None]/self.a**2)*mu).sum()*self.grad_lr[0]

            print('mu: ', self._mu(t)[10].sum())
            print('gradients: ', da, db)
            if (np.abs(da) <= self.grad_conv * np.abs(self.a)) and (np.abs(db) <= self.grad_conv * np.abs(self.b)):
                print('converged: ', iter)
                print('a: ', self.a)
                print('b: ', self.b)
                return;
            self.a = self.a + da
            self.b = self.b + db
            print('a: ', self.a)
            print('b: ', self.b)
        print(self.grad_iter)

    def _grad_ascent_linesearch(self, P_tij):
        t = np.arange(P_tij.shape[0])
        L = lambda : (P_tij*np.log(self.transmat_ + np.transpose(self.A1[:,:,None]*self.U(t),[2,0,1]))).sum()
        # grad_ascent update w/ linesearch
        L_prev = L()
        for iter in range(self.grad_iter):
            # define direction
            mu = self._mu(t) * P_tij
            db = -mu.sum()/self.a
            da = -(((t/24/60/120-self.b)[:,None,None]/self.a**2)*mu).sum()
            delta = np.array([da,db])
            delta =delta/np.linalg.norm(delta, ord=2)
            print('grad: ', da, db)
            # determine step size
            prev_a = self.a
            prev_b = self.b
            updated = False
            for s in self.grad_lr:
                self.a = self.a + s*delta[0]
                self.b = self.b + s*delta[1]
                L_new = L()
                # check if sufficient improvement, if so quit
                #note: log likelihood (L) will always be <0
                #therefore WANT L_new/L_prev < 1
                if L_new/L_prev < self.grad_conv:
                    L_prev = L_new
                    updated = True
                    break;
                else:
                    self.a = prev_a
                    self.b = prev_b

            # check convergence
            # if (np.abs(delta[0]/self.a) < self.grad_conv) and (np.abs(delta[1]/self.b) < self.grad_conv): #TODO: put in convergence criteria
            #     break
            if not updated:
                self.a = prev_a
                self.b = prev_b
                print('not updated')
                break
            self.track_params.append([self.a,self.b, delta[0], delta[1], L_new])
            print('a: ',self.a)
            print('b: ', self.b)

    def _newtons(self,P_tij):
        for i in range(self.grad_iter):
            print(i)
            J, g = self._jacobian_grad(P_tij)
            delta =  - np.matmul(np.linalg.inv(J), g)
            print('Jac: ', J)
            print('grad: ', g)
            print('Diff: ', delta)
            if (np.abs(delta[0]/self.a) < self.grad_conv) and (np.abs(delta[1]/self.b) < self.grad_conv): #TODO: put in convergence criteria
                break
            self.a = self.a + self.grad_lr[0] * delta[0]
            self.b = self.b + self.grad_lr[1] * delta[1]
            self.track_params.append([self.a,self.b, delta[0], delta[1]])
            print('a: ',self.a)
            print('b: ', self.b)
        return

    def _newtons_linesearch(self,P_tij):
        t = np.arange(P_tij.shape[0])
        L = lambda : (P_tij*np.log(self.transmat_ + np.transpose(self.A1[:,:,None]*self.U(t),[2,0,1]))).sum()
        #newton update w/ linesearch
        L_prev = L()
        for i in range(self.grad_iter):
            # get step direction
            print(i)
            J, g = self._jacobian_grad(P_tij)
            delta =  - np.matmul(np.linalg.inv(J), g)
            if np.linalg.norm(delta, ord=2)>1: #sets a maximum to the step size
                delta =delta/np.linalg.norm(delta, ord=2)
            print('Jac: ', J)
            print('grad: ', g)
            print('Diff: ', delta)
            # determine step size
            prev_a = self.a
            prev_b = self.b
            updated = False
            for s in self.grad_lr:
                self.a = self.a + s*delta[0]
                self.b = self.b + s*delta[1]
                L_new = L()
                # check if sufficient improvement, if so quit
                #note: log likelihood (L) will always be <0
                #therefore WANT L_new/L_prev < 1
                print('L',L_new,L_prev)
                if L_new/L_prev < self.grad_conv:
                    L_prev = L_new
                    updated = True
                    break;
                else:
                    self.a = prev_a
                    self.b = prev_b

            # check convergence
            # if (np.abs(delta[0]/self.a) < self.grad_conv) and (np.abs(delta[1]/self.b) < self.grad_conv): #TODO: put in convergence criteria
            #     break
            if not updated:
                self.a = prev_a
                self.b = prev_b
                print('not updated')
                break
            self.track_params.append([self.a,self.b, delta[0], delta[1], L_new])
            print('a: ',self.a)
            print('b: ', self.b)

        return

    def _scipy(self, P_tij):
        t = np.arange(P_tij.shape[0])
        def nL(x):
            a = x[0]
            b = x[1]
            u_loc = lambda t: 1-(1+np.exp(-(t/120/60/24-b)/a))**-1
            t = np.arange(P_tij.shape[0])
            return -(P_tij*np.log(self.transmat_ + np.transpose(self.A1[:,:,None]*u_loc(t),[2,0,1]))).sum()
        def jac(x):
            a = x[0]
            b = x[1]
            sigma_loc = lambda t: 1-(1+np.exp(-(t/120/60/24-b)/a))**-1
            u_loc = lambda t: 1-sigma_loc(t)
            du_loc = lambda t: -sigma_loc(t)*(1-sigma_loc(t))
            mu_loc =lambda t: np.transpose(self.A1[:,:,None]*u_loc(t),[2,0,1])/self.transmat_ + np.transpose(self.A1[:,:,None]*u_loc(t),[2,0,1])
            t = np.arange(P_tij.shape[0])
            mu = mu_loc(t) * P_tij
            dxda = -((t/24/60/120-b)[:,None,None]/a**2)
            dxdb = -1/a
            db = mu*dxdb
            da = mu*dxda
            return -np.array([np.sum(da), np.sum(db)])
        def hess(x):
            a = x[0]
            b = x[1]
            sigma_loc = lambda t: 1-(1+np.exp(-(t/120/60/24-b)/a))**-1
            u_loc = lambda t: 1-sigma_loc(t)
            du_loc = lambda t: -sigma_loc(t)*(1-sigma_loc(t))
            mu_loc =lambda t: np.transpose(self.A1[:,:,None]*u_loc(t),[2,0,1])/self.transmat_ + np.transpose(self.A1[:,:,None]*u_loc(t),[2,0,1])
            t = np.arange(P_tij.shape[0])
            mu = mu_loc(t) * P_tij
            dxda = -((t/24/60/120-b)[:,None,None]/a**2)
            dxdb = -1/a
            db = -mu*dxdb
            da = -mu*dxda
            dL2_dada = (da*(-mu_loc(t)*dxda+(2*u_loc(t)-1)[:,None,None]*dxda-2/a)).sum()
            dL2_dbdb = (db*(-mu_loc(t)*dxdb+(2*u_loc(t)-1)[:,None,None]*dxdb)).sum()
            dL2_dadb = (db*(-mu_loc(t)*dxda+(2*u_loc(t)-1)[:,None,None]*dxda+dxdb)).sum()
            dL2_dbda = dL2_dadb#(db*(-self._mu(t)*dxda+(2*self.U(t)-1)[:,None,None]*dxda+dxdb)
            return np.array([[dL2_dada, dL2_dadb],[dL2_dbda, dL2_dbdb]])

        #result = scipy.optimize.minimize(nL,[self.a,self.b],method='Newton-CG', jac=jac, hess=hess, tol=self.grad_conv,)
        result = scipy.optimize.minimize(nL,[self.a,self.b])
        self.a = result.x[0]
        self.b = result.x[1]
        self.track_params.append([self.a,self.b,result.fun])
        self._grad_ascent_linesearch(P_tij)

    def _jacobian_grad(self,P_tij):
        t = np.arange(P_tij.shape[0])
        mu = self._mu(t) * P_tij
        dxda = -((t/24/60/120-self.b)[:,None,None]/self.a**2)
        dxdb = -1/self.a
        db = mu*dxdb
        da = mu*dxda

        G = np.array([np.sum(da), np.sum(db)])

        dL2_dada = (da*(-self._mu(t)*dxda+(2*self.U(t)-1)[:,None,None]*dxda-2/self.a)).sum()
        dL2_dbdb = (db*(-self._mu(t)*dxdb+(2*self.U(t)-1)[:,None,None]*dxdb)).sum()
        dL2_dadb = (db*(-self._mu(t)*dxda+(2*self.U(t)-1)[:,None,None]*dxda+dxdb)).sum()
        dL2_dbda = dL2_dadb#(db*(-self._mu(t)*dxda+(2*self.U(t)-1)[:,None,None]*dxda+dxdb)
        J = np.array([[dL2_dada, dL2_dadb],[dL2_dbda, dL2_dbdb]])
        return J, G

    # # OLD VERSION: x = a*(t-b)
    # def sigma(self, t):
    #     return (1+np.exp(-self.a*(t-self.b)))**-1
    #
    # def _grad_ascent(self, P_tij):
    #     print('transition posteriors:', P_tij.shape)
    #     print('max: ', np.max(P_tij))
    #     print('sumcheck: ', np.sum(P_tij[1,:,:]), np.sum(P_tij[-1,:,:]) )
    #     t = np.arange(P_tij.shape[0]) + self.t0
    #     for iter in range(self.grad_iter):
    #         mu = self._mu(t) * P_tij
    #         db = -mu.sum()#*self.a*-1
    #         da = ((t[:,None,None]-self.b)*mu).sum()
    #
    #         print('mu: ', self._mu(t)[10].sum())
    #         print('gradients: ', da, db)
    #         if (self.grad_lr * np.abs(da) <= self.grad_conv * np.abs(self.a)) and (self.grad_lr * np.abs(db) <= self.grad_conv * np.abs(self.b)):
    #             print('converged: ', iter)
    #             return;
    #         #self.a = self.a + self.grad_lr * da
    #         self.b = self.b + self.grad_lr * db
    #     print(self.grad_iter)

    def _transition_posterior(self, posteriors):
        print('posterior shape: ',posteriors.shape)
        print('max: ', np.max(posteriors))
        print('sumcheck: ', np.sum(posteriors[1,:]), np.sum(posteriors[-1,:]) )
        P_tij = np.zeros((posteriors.shape[0],self.n_components,self.n_components))
        for t in range(1, P_tij.shape[0]):
            for i in range(self.n_components):
                for j in range(self.n_components):
                    P_tij[t,i,j] = posteriors[t-1,i]*posteriors[t,j]
        return P_tij

    def _mu(self,t):
        num = np.transpose(self.A1[:,:,None]*self.dU(t),[2,0,1])
        denom = self.transmat_ + np.transpose(self.A1[:,:,None]*self.U(t),[2,0,1])
        # print('mu num: ', num[10,:,:].sum())
        # print('mu denom: ', denom[10].sum())
        return num/denom

    def _set_transmat(self, transmat_):
        self.transmat_ = (transmat_)

    def _set_startprob(self,startprob):
        self.startprob_ = startprob