# -*- coding: utf-8 -*-
from __future__ import division
import numpy
import scipy
from numpy import exp, sqrt, pi, sin, cos
from scipy.special import  gammainc, dawsn, erfi
# SciPy no longer has random, and this code needs to be updated. The import is disabled
# for now because it keeps the rest of the tests from running.
# from scipy import random
from scipy import optimize
from math import factorial

class sample:
    def __init__(self, mod, **kwargs):
        r""""
        Sample positions and velocities from DF generated by LIMEPY
        Parameters:
        mod : Python object
          Model generated by LIMEPY
        N : scalar, optional
           Number of points to be generated; default=1000
        seed : scaler, optional
             Random seed
        verbose : bool, optional
                Print diagnostics; default=False
        """

        self._set_params_and_init(mod, **kwargs)
        self._set_masses(mod)
        self._sample_r(mod)
        self._sample_v(mod)
        self._to_cartesian()

        if (self.verbose): self._summary()

    def _set_params_and_init(self, mod, **kwargs):
        if not mod.converged: raise ValueError(" Error: This model has not converged, abort sampling ...")

        if (mod.model=="spes"):
            raise ValueError("Sampling from SPES model not implemented ...")
        self.seed = 199
        self.verbose = False
        self.N = 1000
        if kwargs is not None:
            for key, value in kwargs.items():
                setattr(self, key, value)

        random.seed(self.seed)

        if (mod.multi):
            #  N is recalculated in case of multi mass
            self.Nj = numpy.array(mod.Mj/mod.mj, dtype='int')
            self.N = sum(self.Nj)

            if min(self.Nj)==0: raise ValueError(" Error: One of the mass components has zero stars!")
        self.mod = mod
        self.ani = True if min(mod.raj)/mod.rt < 3 else False

    def _set_masses(self, mod):
        # Create mass array
        if not mod.multi:
            self.m = numpy.zeros(self.N) + mod.M/self.N
            self.ra = mod.ra

        if mod.multi:
            self.Nstart = numpy.zeros(mod.nmbin,dtype='int')
            self.Nend = numpy.zeros(mod.nmbin,dtype='int')
            for j in range(mod.nmbin):
                self.Nstart[j] = int(sum(self.Nj[0:j]))
                self.Nend[j] = int(self.Nstart[j] + self.Nj[j])

            self.m = numpy.zeros(self.N)
            self.ra = numpy.zeros(self.N)
            self.s2j = numpy.zeros(self.N)
            for j in range(mod.nmbin):
                self.m[self.Nstart[j]:self.Nend[j]] = mod.mj[j]
                self.ra[self.Nstart[j]:self.Nend[j]] = mod.raj[j]
                self.s2j[self.Nstart[j]:self.Nend[j]] = mod.s2j[j]
            if (self.verbose):
                print(" N as computed from Mj/mj = ",self.Nj)

    def _sample_r(self, m):
        # Sample radial distances from cumulative mass profile
        if (self.verbose): print(" sample r ...")
        ran = random.rand(self.N)

        if not m.multi:
            self.r = numpy.interp(ran, m.mc/m.mc[-1], m.r)

        if m.multi:
            self.r = numpy.zeros(self.N)
            for j in range(m.nmbin):
                s, e = self.Nstart[j], self.Nend[j]
                m.mcj[j][0] = 0.5*m.mcj[j][1]
                self.r[s:e] = numpy.interp(ran[s:e], m.mcj[j]/m.mcj[j][-1], m.r)

        # get dimensionless potential
        self.phihat = m.interp_phi(self.r)/m.s2

        return

    def _sample_v(self, m):
        # Sample random values for x = k**1.5
        if (self.verbose): print(" sample v ...")
        self.xmax= self.phihat**1.5
        nx = 10
        self.nx = nx
        frac = numpy.linspace(0,1,nx+1)

        self.x = numpy.zeros(len(self.r))
        for j in range(1,nx):
            self.x = numpy.vstack((self.x, self.xmax*frac[j]))
        self.x = numpy.vstack((self.x, self.xmax))

        self.k, self.v, self.rfinal = [], [], []
        for j in range(self.mod.nmbin):
            if (self.verbose): print("   set-up segments for velocity cdf ...")

            self.ycum = numpy.zeros(self.N)
            self.y = self._pdf_k32(self.r, self.phihat, self.x[0], j)
            for i in range(1,nx):
                self.y = numpy.vstack((self.y, self._pdf_k32(self.r, self.phihat, self.x[i], j)))
            self.y = numpy.vstack((self.y, 0*self.r, self.x[-1]))

            if (self.verbose): print("   compute cdf ...")
            self._compute_cdf()

            if (self.verbose): print("   (rejection) sampling of k^3/2 values ...")
            self._sample_k(j)

        self.r = self.rfinal
        self.phihat = m.interp_phi(self.r)/m.s2

        self.phi = -self.phihat*m.s2 - m.G*m.M/m.rt

        self._sample_angles(self.N)

    def _compute_cdf(self):
        for j in range(len(self.x)-1):
            dy = self.y[j]*(self.x[j+1] - self.x[j])
            self.ycum = numpy.vstack((self.ycum, self.ycum[-1]+dy))

    def _Eg(self,x,s):
        return exp(x)*gammainc(s,x) if s>0 else exp(x)

    def _pdf_k32(self, r, f, x, j):
        # PDF of k^3/2
        g = self.mod.g
        x13 = x**(1./3)
        x23 = x**(2./3)
        sig2fac = self.mod.s2/self.mod.s2j[j]
        E = (f-x23)*sig2fac

        c = (x>0)
        func = numpy.zeros(len(x))
        if (sum(c)>0):
            if (self.ani):
                p = r/self.mod.raj[j]
                F = dawsn(x13*p*sqrt(sig2fac))/p
                func[c] = F[c]/(x13[c]*sqrt(sig2fac))*self._Eg(E[c],g)
            else:
                func[c] = self._Eg(E[c],g)
        if (sum(~c)>0): func[~c] = self._Eg(E[~c],g)
        return func

    def _sample_k(self, jbin):
        # y(x) = y_j
        # cdf(x) = ycum_j + y_j*(x - x_j)
        # x = R/y_j + x_j
        N = self.N
        iter=0

        c = numpy.zeros(N, dtype='bool')+True

        if self.mod.multi:
            sta, end = self.Nstart[jbin], self.Nend[jbin]
            c[0:sta] = False
            c[end::] = False

        ycum = self.ycum[:,c]
        rtmp = self.r[c]

        ftmp = self.phihat[c]
        y = self.y[:,c]
        x = self.x[:,c]

        if self.mod.multi:
            c = numpy.zeros(self.Nj[jbin], dtype='bool')+True


        while (sum(c) > 0):
            nc = sum(c)

            ycum = ycum[:,c]
            rtmp = rtmp[c]
            ftmp = ftmp[c]
            y = y[:,c]
            x = x[:,c]

            R = random.rand(nc)*ycum[-1]
            xtmp = numpy.zeros(nc)
            P = numpy.zeros(nc)

            for j in range(self.nx):
                cr = (R>=ycum[j])&(R<ycum[j+1])
                if (sum(cr)>0):
                    xtmp[cr] = (R[cr]-ycum[j,cr])/y[j,cr] + x[j,cr]
                    P[cr] = y[j,cr]

            R2 = random.rand(nc)*P
            f = self._pdf_k32(rtmp, ftmp, xtmp, jbin)

            c = (R2>f)

            if (iter==0):
                xsamp = xtmp[~c]
                r = rtmp[~c]
            else:
                xsamp = numpy.r_[xsamp, xtmp[~c]]
                r = numpy.r_[r, rtmp[~c]]
            iter+=1

        # Assign final velocities and update r because of shuffling
        self.k = numpy.r_[self.k, xsamp**(2./3)]

        self.v = numpy.r_[self.v, sqrt(2*xsamp**(2./3)*self.mod.s2)]

        self.rfinal = numpy.r_[self.rfinal, r]

    def _pdf_angle(self, q, *args):
        # Sample random values for: q = cos(theta)
        # P(q) = erfi(sqrt(k)*p*q)/erfi(sqrt(k)*p)

        a, R = args
        return R  - erfi(a*q)/erfi(a)

    def _sample_angles(self,N):
        if (self.verbose): print("   sample angles ...")
        R = random.rand(N)

        if self.ani:
            # Anisotropic systems sample angles: cdf(q) = erfi(a*q)/erfi(a), a = sqrt(k)*p
            p = self.r/self.ra
            a = p*sqrt(self.k)
            if self.mod.multi:
                a *= sqrt(self.mod.s2/self.s2j)
            self.q = numpy.zeros(N)

            for j in range(N):
                self.q[j] = optimize.brentq(self._pdf_angle, 0, 1, args=(a[j], R[j]))
        else:
            # Isotropic: cdf(q) = q
            self.q = R

        self.vr = self.v*self.q*random.choice((-1,1),size=self.N)
        self.vt = self.v*sqrt(1-self.q**2)

    def _to_cartesian(self):
        if (self.verbose): print(" convert to cartesian coordinates ...")
        r = self.r
        r2 = r**2
        R1 = random.rand(self.N)
        R2 = random.rand(self.N)
        self.x = (1-2*R1)*self.r
        self.y = sqrt(r2 - self.x**2)*cos(2*pi*R2)
        self.z = sqrt(r2 - self.x**2)*sin(2*pi*R2)
        if (self.ani):
            R1 = random.rand(self.N)
            vphi = self.vt*cos(2*pi*R1)
            vtheta = self.vt*sin(2*pi*R1)
            theta = numpy.arccos(self.z/r)
            phi = numpy.arctan2(self.y, self.x)

            self.vx = self.vr*sin(theta)*cos(phi) + vtheta*cos(theta)*cos(phi) - vphi*sin(phi)
            self.vy = self.vr*sin(theta)*sin(phi) + vtheta*cos(theta)*sin(phi) + vphi*cos(phi)
            self.vz = self.vr*cos(theta) - vtheta*sin(theta)
        else:
            v2 = self.v**2
            R1 = random.rand(self.N)
            R2 = random.rand(self.N)
            self.vx = (1-2*R1)*self.v
            self.vy = sqrt(v2 - self.vx**2)*cos(2*pi*R2)
            self.vz = sqrt(v2 - self.vx**2)*sin(2*pi*R2)

    def _summary(self):
        print(" done! ")
        m = self.mod
        f = -self.phihat*m.s2 - m.G*m.M/m.rt

        print("       U: sample = %12.4e; model = %12.4e"%(0.5*sum(self.m*f), m.U))
        print("       K: sample = %12.4e; model = %12.4e"%(sum(0.5*self.m*self.v**2), m.K))
        print("       Q: sample = %12.4e; model = %12.4e"%(sum(0.5*self.m*self.v**2)/(0.5*sum(self.m*f)), m.K/m.U))
        print("  2Kr/Kt: sample = %12.4f; model = %12.4f"%(2*sum(self.vr**2)/sum(self.vt**2),2*m.Kr/m.Kt))

        if m.multi:

            print("\n  Components: ")
            for j in range(m.nmbin):
                c=(self.m==m.mj[j])
                print("      Kj: sample = %12.4e; model = %12.4e"%(sum(0.5*self.m[c]*self.v[c]**2), m.Kj[j]))

            for j in range(m.nmbin):
                c=(self.m==m.mj[j])
                print(" 2Kr/Ktj: sample = %12.4e; model = %12.4e"%(2*sum(self.vr[c]**2)/sum(self.vt[c]**2),2*m.Krj[j]/m.Ktj[j]))


