# copyright ############################### #
# This file is part of the Xpart Package.   #
# Copyright (c) CERN, 2021.                 #
# ######################################### #

import logging

import numpy as np

import xobjects as xo
from .general import _print

import xtrack as xt

logger = logging.getLogger(__name__)

def _check_lengths(**kwargs):
    length = None
    for nn, xx in kwargs.items():
        if hasattr(xx, "__iter__"):
            if hasattr(xx, 'shape') and len(xx.shape) == 0:
                continue
            if length is None:
                length = len(xx)
            else:
                if length != len(xx):
                    raise ValueError(f"invalid length len({nn})={len(xx)}")
    if 'num_particles' in kwargs.keys():
        num_particles = kwargs['num_particles']
        if num_particles is not None and length is None:
            length = num_particles
        if num_particles is not None and length != num_particles:
            raise ValueError(
              f"num_particles={num_particles} is inconsistent with array length")

    if length is None:
        length = 1
    return length

def build_particles(_context=None, _buffer=None, _offset=None, _capacity=None,
                      mode=None,
                      particle_ref=None,
                      num_particles=None,
                      x=None, px=None, y=None, py=None,
                      zeta=None, delta=None, pzeta=None, ptau=None,
                      x_norm=None, px_norm=None, y_norm=None, py_norm=None,
                      zeta_norm=None, pzeta_norm=None,
                      tracker=None,
                      line=None,
                      at_element=None,
                      match_at_s=None,
                      particle_on_co=None,
                      R_matrix=None,
                      W_matrix=None,
                      method=None,
                      nemitt_x=None, nemitt_y=None,nemitt_zeta = None,
                      scale_with_transverse_norm_emitt=None,
                      weight=None,
                      s_tol=1e-6,
                      include_collective=False,
                      **kwargs, # They are passed to the twiss
                    ):

    """
    Same as `xtrack.Line.build_particles`. See there for documentation.

    """

    if line is not None and tracker is not None:
        raise ValueError(
            'line and tracker cannot be provided at the same time.')

    if tracker is not None:
        _print('Warning! '
            "The argument tracker is deprecated. Please use line instead.")
        line = tracker.line

    if line is not None:
        assert line.tracker is not None, ("The line must have a tracker, "
            "please call Line.build_tracker() first.")

    assert mode in [None, 'set', 'shift', 'normalized_transverse']
    Particles = xt.Particles  # To get the right Particles class depending on pyheatail interface state

    assert 'at_s' not in kwargs, "at_s is not a valid argument for this function"

    # Deprecation warning for scale_with_transverse_norm_emitt
    if scale_with_transverse_norm_emitt is not None:
        _print('Warning! '
            "scale_with_transverse_norm_emitt is deprecated. Use "
            "nemitt_x and nemitt_y instead.")

    if (particle_ref is not None and particle_on_co is not None):
        raise ValueError("`particle_ref` and `particle_on_co`"
                " cannot be provided at the same time")

    if particle_on_co is None and particle_ref is None:
        if line is not None:
            particle_ref = line.particle_ref

    if particle_ref is None:
        assert particle_on_co is not None, (
            "`particle_ref` or `particle_on_co` must be provided!")
        particle_ref = particle_on_co

    if not isinstance(particle_ref._buffer.context, xo.ContextCpu):
        particle_ref = particle_ref.copy(_context=xo.ContextCpu())

    # Move other input parameters to cpu if needed
    # Generated by:
    # for nn in 'x px y py zeta delta x_norm px_norm y_norm py_norm'.split():
    #     _print(f'{nn} = ({nn}.get() if hasattr({nn}, "get") else {nn})')
    x = (x.get() if hasattr(x, "get") else x)
    px = (px.get() if hasattr(px, "get") else px)
    y = (y.get() if hasattr(y, "get") else y)
    py = (py.get() if hasattr(py, "get") else py)
    zeta = (zeta.get() if hasattr(zeta, "get") else zeta)
    delta = (delta.get() if hasattr(delta, "get") else delta)
    pzeta = (pzeta.get() if hasattr(pzeta, "get") else pzeta)
    ptau = (ptau.get() if hasattr(ptau, "get") else ptau)
    x_norm = (x_norm.get() if hasattr(x_norm, "get") else x_norm)
    px_norm = (px_norm.get() if hasattr(px_norm, "get") else px_norm)
    y_norm = (y_norm.get() if hasattr(y_norm, "get") else y_norm)
    py_norm = (py_norm.get() if hasattr(py_norm, "get") else py_norm)
    zeta_norm = (zeta_norm.get() if hasattr(zeta_norm, "get") else zeta_norm)
    pzeta_norm = (pzeta_norm.get() if hasattr(pzeta_norm, "get") else pzeta_norm)

    if line is not None and line.iscollective and not include_collective:
        logger.warning('Ignoring collective elements in particles generation.')
        line = line._get_non_collective_line()

    # Compute pzeta from delta
    if delta is not None:
        assert pzeta is None
        assert ptau is None
        if not np.isscalar(delta):
            delta = np.array(delta)
        beta0 = particle_ref._xobject.beta0[0]
        delta_beta0 = delta * beta0
        ptau_beta0 = (delta_beta0 * delta_beta0
                            + 2. * delta_beta0 * beta0 + 1.)**0.5 - 1.
        pzeta = ptau_beta0 / beta0 / beta0

    # Compute pzeta from ptau
    if ptau is not None:
        assert pzeta is None
        assert delta is None
        if not np.isscalar(ptau):
            ptau = np.array(ptau)
        pzeta = ptau / particle_ref._xobject.beta0[0]

    if (x_norm is not None or px_norm is not None
            or y_norm is not None or py_norm is not None
            or zeta_norm is not None or pzeta_norm is not None):

        if mode is None:
            mode = 'normalized_transverse'
        else:
            assert mode == 'normalized_transverse'

    if mode is None:
        mode = 'set'

    assert particle_ref._capacity == 1
    ref_dict = {
        'q0': particle_ref.q0,
        'mass0': particle_ref.mass0,
        't_sim': particle_ref.t_sim,
        'p0c': particle_ref.p0c[0],
        'gamma0': particle_ref.gamma0[0],
        'beta0': particle_ref.beta0[0],
        'pdg_id': particle_ref.pdg_id[0],
        'anomalous_magnetic_moment': particle_ref.anomalous_magnetic_moment[0],
    }
    part_dict = ref_dict.copy()

    if at_element is not None or match_at_s is not None:
        # Only this case is covered if not starting at element 0
        assert line is not None
        assert mode == 'normalized_transverse'
        if isinstance(at_element, str):
            at_element = line._element_names_unique.index(at_element)
        assert R_matrix is None # Not clear if it is at the element or at start machine
        if particle_on_co is not None:
            assert particle_on_co._xobject.at_element == 0

    if match_at_s is not None:
        assert at_element is not None, (
            'If `match_at_s` is provided, `at_element` needs to be provided and'
            'needs to correspond to the corresponding element in the sequence'
        )
        s_elements = line.get_s_elements()
        s_at_element = s_elements[at_element]
        if np.abs(match_at_s - s_at_element) < s_tol:
            at_element_line_rmat = at_element
            line_rmat = line
            match_at_s = None
        else:
            # Match at a position where there is no marker and backtrack to the previous marker
            expected_at_element = np.where(np.array(s_elements)<=match_at_s)[0][-1]
            assert at_element == expected_at_element or (
                    at_element < expected_at_element and
                        all([xt._is_aperture(line.element_dict[nn], line)
                            or xt._behaves_like_drift(line.element_dict[nn], line)
                    for nn in line._element_names_unique[at_element:expected_at_element]])), (
                "`match_at_s` can only be placed in the drifts downstream of the "
                "specified `at_element`. No active element can be present in between."
                )
            (tracker_rmat, _
                ) = xt.twiss._build_auxiliary_tracker_with_extra_markers(
                    tracker=line.tracker, at_s=[match_at_s],
                    marker_prefix='xpart_rmat_')
            at_element_line_rmat = tracker_rmat.line._element_names_unique.index(
                                                                    'xpart_rmat_0')
            line_rmat = tracker_rmat.line
    else:
        line_rmat = line
        at_element_line_rmat = at_element

    if mode == 'normalized_transverse':

        if W_matrix is None and line is not None:
            if method is not None:
                kwargs['method'] = method
            tw = line_rmat.twiss(particle_on_co=particle_on_co,
                                    particle_ref=particle_ref,
                                    R_matrix=R_matrix, **kwargs)
            tw_state = tw.get_twiss_init(at_element=
                (at_element_line_rmat if at_element_line_rmat is not None else 0))

            # This is not initialized by get_twiss_init
            tw_state.particle_on_co.at_element = line_rmat._element_names_unique.index(
                                                        tw_state.element_name)

            WW = tw_state.W_matrix
            particle_on_co = tw_state.particle_on_co
        elif W_matrix is None and R_matrix is not None:
            import xtrack.linear_normal_form as lnf
            WW, _, _, _ = lnf.compute_linear_normal_form(R_matrix, **kwargs)
        else:
            WW = W_matrix

        num_particles = _check_lengths(num_particles=num_particles,
            x=x, px=px, y=y, py=py, zeta=zeta, pzeta=pzeta,
            x_norm=x_norm, px_norm=px_norm,
            y_norm=y_norm, py_norm=py_norm,
            zeta_norm=zeta_norm, pzeta_norm=pzeta_norm)

        if scale_with_transverse_norm_emitt is not None:
            assert len(scale_with_transverse_norm_emitt) == 2
            assert nemitt_x is None and nemitt_y is None, (
                "If `scale_with_transverse_norm_emitt` is provided, "
                "`nemitt_x` and `nemitt_y` should not be provided.")

            nemitt_x = scale_with_transverse_norm_emitt[0]
            nemitt_y = scale_with_transverse_norm_emitt[1]

        if nemitt_x is None:
            gemitt_x = 1
        else:
            gemitt_x = (nemitt_x / particle_ref._xobject.beta0[0]
                        / particle_ref._xobject.gamma0[0])

        if nemitt_y is None:
            gemitt_y = 1
        else:
            gemitt_y = (nemitt_y / particle_ref._xobject.beta0[0]
                        / particle_ref._xobject.gamma0[0])

        if nemitt_zeta is None:
            gemitt_zeta = 1
        else:
            gemitt_zeta = (nemitt_zeta / particle_ref._xobject.beta0[0]
                        / particle_ref._xobject.gamma0[0])


        n_constraints = sum([vv is not None for vv in [x, x_norm, px, px_norm,
            y, y_norm, py, py_norm, zeta, zeta_norm, pzeta, pzeta_norm]])

        if n_constraints > 6:
            raise ValueError('Too many constraints. At most 6 among `x`, `x_norm`, '
                                '`px`, `px_norm`, `y`, `y_norm`, `py`, `py_norm`, '
                                '`zeta`, `zeta_norm`, `pzeta`, `pzeta_norm` can be '
                                'provided.')
        while n_constraints < 6:
            if x is None and x_norm is None:
                x_norm = 0
                n_constraints += 1
                continue
            if px is None and px_norm is None:
                px_norm = 0
                n_constraints += 1
                continue
            if y is None and y_norm is None:
                y_norm = 0
                n_constraints += 1
                continue
            if py is None and py_norm is None:
                py_norm = 0
                n_constraints += 1
                continue
            if zeta is None and zeta_norm is None:
                zeta_norm = 0
                n_constraints += 1
                continue
            if pzeta is None and pzeta_norm is None:
                pzeta_norm = 0
                n_constraints += 1
                continue

        # if sum([vv is not None for vv in [x, x_norm, px, px_norm]]) <= 1:
        #     if x is None and x_norm is None:
        #         x_norm = 0
        #     if px is None and px_norm is None:
        #         px_norm = 0

        # if sum([vv is not None for vv in [y, y_norm, py, py_norm]]) <= 1:
        #     if y is None and y_norm is None:
        #         y_norm = 0
        #     if py is None and py_norm is None:
        #         py_norm = 0

        # if sum([vv is not None for vv in [zeta, zeta_norm, pzeta, pzeta_norm]]) <= 1:
        #     if zeta is None and zeta_norm is None:
        #         zeta_norm = 0
        #     if pzeta is None and pzeta_norm is None:
        #         pzeta_norm = 0

        BB = np.zeros(shape=(12, num_particles), dtype=np.float64)
        AA = np.zeros(shape=(12, 12), dtype=np.float64)

        # The first 6 equations are X - WW * X_norm = X_CO
        AA[:6, :6] = np.eye(6)
        AA[:6, 6:] = -WW

        BB[0, :] = particle_on_co._xobject.x[0]
        BB[1, :] = particle_on_co._xobject.px[0]
        BB[2, :] = particle_on_co._xobject.y[0]
        BB[3, :] = particle_on_co._xobject.py[0]
        BB[4, :] = particle_on_co._xobject.zeta[0]
        BB[5, :] = particle_on_co._xobject.ptau[0] / particle_on_co._xobject.beta0[0]

        # The next 6 equations fix either X or X_norm
        i_fill = 6
        for ii, rr in enumerate([x, px, y, py, zeta, pzeta]):
            if rr is not None:
                BB[i_fill, :] = np.array(rr)
                AA[i_fill, ii] = 1
                i_fill += 1

        for ii, (rr_norm,  gemitt) in enumerate(zip(
                    [x_norm, px_norm, y_norm, py_norm, zeta_norm, pzeta_norm],
                    [gemitt_x, gemitt_x, gemitt_y, gemitt_y, gemitt_zeta, gemitt_zeta])):
            if rr_norm is not None:
                rr_norm_scaled = np.sqrt(gemitt) * np.array(rr_norm)
                BB[i_fill, :] = rr_norm_scaled
                AA[i_fill, ii + 6] = 1
                i_fill += 1

        # Solve and extract geometric coordinates
        X_Xhat = np.linalg.solve(AA, BB)
        XX = X_Xhat[:6, :]

    elif mode == 'set':

        if x is None: x = 0
        if px is None: px = 0
        if y is None: y = 0
        if py is None: py = 0
        if zeta is None: zeta = 0
        if pzeta is None: pzeta = 0

        if R_matrix is not None:
            logger.warning('R_matrix provided but not used in this mode!')

        num_particles = _check_lengths(num_particles=num_particles,
            zeta=zeta, delta=delta, x=x, px=px,
            y=y, py=py)

        XX = np.zeros(shape=(6, num_particles), dtype=np.float64)
        XX[0, :] = x
        XX[1, :] = px
        XX[2, :] = y
        XX[3, :] = py
        XX[4, :] = zeta
        XX[5, :] = pzeta

    elif mode == "shift":

        if x is None: x = 0
        if px is None: px = 0
        if y is None: y = 0
        if py is None: py = 0
        if zeta is None: zeta = 0
        if pzeta is None: pzeta = 0

        if R_matrix is not None:
            logger.warning('R_matrix provided but not used in this mode!')

        num_particles = _check_lengths(num_particles=num_particles,
            zeta=zeta, delta=delta, x=x, px=px,
            y=y, py=py)

        XX = np.zeros(shape=(6, num_particles), dtype=np.float64)
        XX[0, :] = x + particle_ref.x
        XX[1, :] = px + particle_ref.px
        XX[2, :] = y + particle_ref.y
        XX[3, :] = py + particle_ref.py
        XX[4, :] = zeta + particle_ref.zeta
        XX[5, :] = pzeta + particle_ref.ptau / particle_ref._xobject.beta0[0]
    else:
        raise ValueError('What?!')

    part_dict['x'] = XX[0, :]
    part_dict['px'] = XX[1, :]
    part_dict['y'] = XX[2, :]
    part_dict['py'] = XX[3, :]
    part_dict['zeta'] = XX[4, :]
    part_dict['ptau'] = XX[5, :] * particle_ref._xobject.beta0[0]

    part_dict['weight'] = np.ones(num_particles, dtype=np.float64)

    if _context is None and _buffer is None and line is not None:
        _context = line._buffer.context

    particles = Particles(_context=_context, _buffer=_buffer, _offset=_offset,
                          _capacity=_capacity,**part_dict)

    particles.particle_id[:num_particles] = particles._buffer.context.nparray_to_context_array(
                                   np.arange(0, num_particles, dtype=np.int64))
    if weight is not None:
        particles.weight[:num_particles] = weight

    if match_at_s is not None:
        # Backtrack to at_element
        length_aux_drift = -match_at_s + line.get_s_position(at_element)
        assert length_aux_drift <= 0
        auxdrift = xt.Drift(length=length_aux_drift,
                            _context=line._buffer.context)
        auxdrift.track(particles)

    if at_element is not None:
        if match_at_s is not None:
            particles.s[:num_particles] = particle_on_co._xobject.s[0] + length_aux_drift
        else:
            assert particle_on_co.at_element[0] == at_element
            particles.s[:num_particles] = particle_on_co._xobject.s[0]
        particles.at_element[:num_particles] = at_element

        particles.start_tracking_at_element = at_element

    particles.spin_x[:num_particles] = particle_ref._xobject.spin_x[0]
    particles.spin_y[:num_particles] = particle_ref._xobject.spin_y[0]
    particles.spin_z[:num_particles] = particle_ref._xobject.spin_z[0]
    if 'spin_x' in kwargs:
        particles.spin_x[:num_particles] = kwargs['spin_x']
    if 'spin_y' in kwargs:
        particles.spin_y[:num_particles] = kwargs['spin_y']
    if 'spin_z' in kwargs:
        particles.spin_z[:num_particles] = kwargs['spin_z']

    return particles
