import sys
import os
import glob
import numpy
import time

from PyMca5.PyMca import specfilewrapper, EdfFile, SixCircle, specfile

from .. import backend, errors, util


class pixels(backend.ProjectionBase):
    def project(self, wavelength, UB, gamma, delta, theta, mu, chi, phi):
        y, x = numpy.mgrid[slice(None, gamma.shape[0]), slice(None, delta.shape[0])]
        return (y, x)

    def get_axis_labels(self):
        return "y", "x"


class HKLProjection(backend.ProjectionBase):
    # arrays: gamma, delta
    # scalars: theta, mu, chi, phi
    def project(self, wavelength, UB, beta, delta, omega, alfa, chi, phi):
        R = SixCircle.getHKL(
            wavelength,
            UB,
            gamma=beta,
            delta=delta,
            theta=omega,
            mu=alfa,
            chi=chi,
            phi=phi,
        )
        shape = beta.size, delta.size
        H = R[0, :].reshape(shape)
        K = R[1, :].reshape(shape)
        L = R[2, :].reshape(shape)
        return (H, K, L)

    def get_axis_labels(self):
        return "H", "K", "L"


class HKProjection(HKLProjection):
    def project(self, wavelength, UB, beta, delta, omega, alfa, chi, phi):
        H, K, L = super(HKProjection, self).project(
            wavelength, UB, beta, delta, omega, alfa, chi, phi
        )
        return (H, K)

    def get_axis_labels(self):
        return "H", "K"


class ThetaLProjection(backend.ProjectionBase):
    # arrays: gamma, delta
    # scalars: theta, mu, chi, phi
    def project(self, wavelength, UB, gamma, delta, theta, mu, chi, phi):
        R = SixCircle.getHKL(
            wavelength,
            UB,
            gamma=gamma,
            delta=delta,
            theta=theta,
            mu=mu,
            chi=chi,
            phi=phi,
        )
        shape = gamma.size, delta.size
        L = R[2, :].reshape(shape)
        theta_array = numpy.ones_like(L) * theta
        return (theta_array, L)

    def get_axis_labels(self):
        return "Theta", "L"


class QProjection(backend.ProjectionBase):
    def project(self, wavelength, UB, gamma, delta, theta, mu, chi, phi):
        shape = gamma.size, delta.size
        sixc = SixCircle.SixCircle()
        sixc.setLambda(wavelength)
        sixc.setUB(UB)
        R = sixc.getQSurface(
            gamma=gamma, delta=delta, theta=theta, mu=mu, chi=chi, phi=phi
        )
        qz = R[0, :].reshape(shape)
        qy = R[1, :].reshape(shape)
        qx = R[2, :].reshape(shape)
        return (qz, qy, qx)

    def get_axis_labels(self):
        return "qx", "qy", "qz"


class SphericalQProjection(QProjection):
    def project(self, wavelength, UB, gamma, delta, theta, mu, chi, phi):
        qz, qy, qx = super(SphericalQProjection, self).project(
            wavelength, UB, gamma, delta, theta, mu, chi, phi
        )
        q = numpy.sqrt(qx ** 2 + qy ** 2 + qz ** 2)
        theta = numpy.arccos(qz / q)
        phi = numpy.arctan2(qy, qx)
        return (q, theta, phi)

    def get_axis_labels(self):
        return "Q", "Theta", "Phi"


class CylindricalQProjection(QProjection):
    def project(self, wavelength, UB, gamma, delta, theta, mu, chi, phi):
        qz, qy, qx = super(CylindricalQProjection, self).project(
            wavelength, UB, gamma, delta, theta, mu, chi, phi
        )
        qpar = numpy.sqrt(qx ** 2 + qy ** 2)
        phi = numpy.arctan2(qy, qx)
        return (qpar, qz, phi)

    def get_axis_labels(self):
        return "qpar", "qz", "Phi"


class nrQProjection(backend.ProjectionBase):
    def project(self, wavelength, UB, gamma, delta, theta, mu, chi, phi):
        k0 = 2 * numpy.pi / wavelength
        delta, gamma = numpy.meshgrid(delta, gamma)
        mu *= numpy.pi / 180
        delta *= numpy.pi / 180
        gamma *= numpy.pi / 180

        qy = k0 * (
            numpy.cos(gamma) * numpy.cos(delta) - numpy.cos(mu)
        )  ## definition of qx, and qy same as spec at theta = 0
        qx = k0 * (numpy.cos(gamma) * numpy.sin(delta))
        qz = k0 * (numpy.sin(gamma) + numpy.sin(mu))
        return (qx, qy, qz)

    def get_axis_labels(self):
        return "qx", "qy", "qz"


class TwoThetaProjection(SphericalQProjection):
    def project(self, wavelength, UB, gamma, delta, theta, mu, chi, phi):
        q, theta, phi = super(TwoThetaProjection, self).project(
            wavelength, UB, gamma, delta, theta, mu, chi, phi
        )
        return (
            2 * numpy.arcsin(q * wavelength / (4 * numpy.pi)) / numpy.pi * 180,
        )  # note: we need to return a 1-tuple?

    def get_axis_labels(self):
        return "TwoTheta"


class Qpp(nrQProjection):
    def project(self, wavelength, UB, gamma, delta, theta, mu, chi, phi):
        qx, qy, qz = super(Qpp, self).project(
            wavelength, UB, gamma, delta, theta, mu, chi, phi
        )
        qpar = numpy.sqrt(qx ** 2 + qy ** 2)
        qpar[numpy.sign(qx) == -1] *= -1
        return (qpar, qz)

    def get_axis_labels(self):
        return "Qpar", "Qz"


class GammaDeltaTheta(
    HKLProjection
):  # just passing on the coordinates, makes it easy to accurately test the theta correction
    def project(self, wavelength, UB, gamma, delta, theta, mu, chi, phi):
        delta, gamma = numpy.meshgrid(delta, gamma)
        theta = theta * numpy.ones_like(delta)
        return (gamma, delta, theta)

    def get_axis_labels(self):
        return "Gamma", "Delta", "Theta"


class GammaDelta(
    HKLProjection
):  # just passing on the coordinates, makes it easy to accurately test the theta correction
    def project(self, wavelength, UB, gamma, delta, theta, mu, chi, phi):
        delta, gamma = numpy.meshgrid(delta, gamma)
        return (gamma, delta)

    def get_axis_labels(self):
        return "Gamma", "Delta"


class GammaDeltaMu(
    HKLProjection
):  # just passing on the coordinates, makes it easy to accurately test the theta correction
    def project(self, wavelength, UB, gamma, delta, theta, mu, chi, phi):
        delta, gamma = numpy.meshgrid(delta, gamma)
        mu = mu * numpy.ones_like(delta)
        return (gamma, delta, mu)

    def get_axis_labels(self):
        return "Gamma", "Delta", "Mu"


class BM32Input(backend.InputBase):
    # OFFICIAL API

    dbg_scanno = None
    dbg_pointno = None

    def generate_jobs(self, command):
        scans = util.parse_multi_range(",".join(command).replace(" ", ","))
        if not len(scans):
            sys.stderr.write("error: no scans selected, nothing to do\n")
        for scanno in scans:
            util.status("processing scan {0}...".format(scanno))
            scan = self.get_scan(scanno)
            if self.config.pr:
                pointcount = self.config.pr[1] - self.config.pr[0] + 1
                start = self.config.pr[0]
            else:
                start = 0
                try:
                    pointcount = scan.lines()
                except specfile.error:  # no points
                    continue
            next(self.get_images(scan, 0, pointcount - 1, dry_run=True))  # dryrun
            if pointcount > self.config.target_weight * 1.4:
                for s in util.chunk_slicer(pointcount, self.config.target_weight):
                    yield backend.Job(
                        scan=scanno,
                        firstpoint=start + s.start,
                        lastpoint=start + s.stop - 1,
                        weight=s.stop - s.start,
                    )
            else:
                yield backend.Job(
                    scan=scanno,
                    firstpoint=start,
                    lastpoint=start + pointcount - 1,
                    weight=pointcount,
                )

    def process_job(self, job):
        super(BM32Input, self).process_job(job)
        scan = self.get_scan(job.scan)
        self.metadict = dict()
        try:
            scanparams = self.get_scan_params(scan)  # wavelength, UB
            pointparams = self.get_point_params(
                scan, job.firstpoint, job.lastpoint
            )  # 2D array of diffractometer angles + mon + transm
            images = self.get_images(scan, job.firstpoint, job.lastpoint)  # iterator!

            for pp, image in zip(pointparams, images):
                yield self.process_image(scanparams, pp, image)
            util.statuseol()
        except Exception as exc:
            # exc.args = errors.addmessage(exc.args, ', An error occured for scan {0} at point {1}. See above for more information'.format(self.dbg_scanno, self.dbg_pointno))
            raise
        self.metadata.add_section("id03_backend", self.metadict)

    def parse_config(self, config):
        super(BM32Input, self).parse_config(config)
        self.config.xmask = util.parse_multi_range(
            config.pop("xmask", None)
        )  # Optional, select a subset of the image range in the x direction. all by default
        self.config.ymask = util.parse_multi_range(
            config.pop("ymask", None)
        )  # Optional, select a subset of the image range in the y direction. all by default
        self.config.specfile = config.pop("specfile")  # Location of the specfile
        self.config.imagefolder = config.pop(
            "imagefolder", None
        )  # Optional, takes specfile folder tag by default
        self.config.pr = util.parse_tuple(
            config.pop("pr", None), length=2, type=int
        )  # Optional, all range by default
        self.config.background = config.pop(
            "background", None
        )  # Optional, if supplied a space of this image is constructed
        if self.config.xmask is None:
            self.config.xmask = slice(None)
        if self.config.ymask is None:
            self.config.ymask = slice(None)
        self.config.maskmatrix = load_matrix(
            config.pop("maskmatrix", None)
        )  # Optional, if supplied pixels where the mask is 0 will be removed
        self.config.sdd = config.pop("sdd", None)  # sample to detector distance (mm)
        if self.config.sdd is not None:
            self.config.sdd = float(self.config.sdd)
        self.config.pixelsize = util.parse_tuple(
            config.pop("pixelsize", None), length=2, type=float
        )  # pixel size x/y (mm) (same dimension as sdd)

    def get_destination_options(self, command):
        if not command:
            return False
        command = ",".join(command).replace(" ", ",")
        scans = util.parse_multi_range(command)
        return dict(
            first=min(scans),
            last=max(scans),
            range=",".join(str(scan) for scan in scans),
        )

    # CONVENIENCE FUNCTIONS
    _spec = None

    def get_scan(self, scannumber):
        if self._spec is None:
            self._spec = specfilewrapper.Specfile(self.config.specfile)
        return self._spec.select("{0}.1".format(scannumber))

    def find_edfs(self, pattern):
        files = glob.glob(pattern)
        ret = {}
        for file in files:
            try:
                filename = os.path.basename(file).split(".")[0]
                imno = int(filename.split("_")[-1].split("-")[-1])
                ret[imno] = file
            except ValueError:
                continue
        return ret

    @staticmethod
    def apply_mask(data, xmask, ymask):
        roi = data[ymask, :]
        return roi[:, xmask]

    # MAIN LOGIC
    def get_scan_params(self, scan):
        self.dbg_scanno = scan.number()
        UB = numpy.array(scan.header("G")[2].split(" ")[-9:], dtype=numpy.float)
        wavelength = float(scan.header("G")[1].split(" ")[-1])

        self.metadict["UB"] = UB
        self.metadict["wavelength"] = wavelength

        return wavelength, UB

    def get_images(self, scan, first, last, dry_run=False):
        imagenos = (
            numpy.array(scan.datacol("img")[slice(first, last + 1)], dtype=numpy.int)
            + 1
        )  ##error in spec?!
        if self.config.background:
            if not os.path.exists(self.config.background):
                raise errors.FileError(
                    "could not find background file {0}".format(self.config.background)
                )
            if dry_run:
                yield
            else:
                edf = EdfFile.EdfFile(self.config.background)
                for i in range(first, last + 1):
                    self.dbg_pointno = i
                    yield edf
        else:
            try:
                uccdtagline = scan.header("M")[0].split()[-1]
                UCCD = os.path.dirname(uccdtagline).split(os.sep)
            except:
                print(
                    "warning: UCCD tag not found, use imagefolder for proper file specification"
                )
                UCCD = []
            pattern = self._get_pattern(UCCD)
            matches = self.find_edfs(pattern)
            if not set(imagenos).issubset(set(matches.keys())):
                raise errors.FileError(
                    "incorrect number of matches for scan {0} using pattern {1}".format(
                        scan.number(), pattern
                    )
                )
            if dry_run:
                yield
            else:
                for i in imagenos:
                    self.dbg_pointno = i
                    edf = EdfFile.EdfFile(matches[i])
                    yield edf

    def _get_pattern(self, UCCD):
        imagefolder = self.config.imagefolder
        if imagefolder:
            try:
                imagefolder = imagefolder.format(UCCD=UCCD, rUCCD=list(reversed(UCCD)))
            except Exception as e:
                raise errors.ConfigError(
                    "invalid 'imagefolder' specification '{0}': {1}".format(
                        self.config.imagefolder, e
                    )
                )
            else:
                if not os.path.exists(imagefolder):
                    raise errors.ConfigError(
                        "invalid 'imagefolder' specification '{0}'. Path {1} does not exist".format(
                            self.config.imagefolder, imagefolder
                        )
                    )
        else:
            imagefolder = os.path.join(*UCCD)
            if not os.path.exists(imagefolder):
                raise errors.ConfigError(
                    "invalid UCCD tag '{0}'. The UCCD tag in the specfile does not point to an existing folder. Specify the imagefolder in the configuration file.".format(
                        imagefolder
                    )
                )
        return os.path.join(imagefolder, "*")


class EH1(BM32Input):
    def parse_config(self, config):
        super(EH1, self).parse_config(config)
        self.config.centralpixel = util.parse_tuple(
            config.pop("centralpixel", None), length=2, type=int
        )
        self.config.UB = util.parse_tuple(config.pop("ub", None), length=9, type=float)

    def process_image(self, scanparams, pointparams, edf):
        delta, omega, alfa, beta, chi, phi, mon, transm = pointparams
        wavelength, UB = scanparams

        image = edf.GetData(0)
        header = edf.GetHeader(0)

        weights = numpy.ones_like(image)

        if not self.config.centralpixel:
            self.config.centralpixel = (int(header["y_beam"]), int(header["x_beam"]))
        if not self.config.sdd:
            self.config.sdd = float(header["det_sample_dist"])

        if self.config.background:
            data = image / mon
        else:
            data = image / mon / transm

        if mon == 0:
            raise errors.BackendError(
                "Monitor is zero, this results in empty output. Scannumber = {0}, pointnumber = {1}. Did you forget to open the shutter?".format(
                    self.dbg_scanno, self.dbg_pointno
                )
            )

        util.status(
            "{4}| beta: {0:.3f}, delta: {1:.3f}, omega: {2:.3f}, alfa: {3:.3f}".format(
                beta, delta, omega, alfa, time.ctime(time.time())
            )
        )

        # pixels to angles
        pixelsize = numpy.array(self.config.pixelsize)
        sdd = self.config.sdd

        app = numpy.arctan(pixelsize / sdd) * 180 / numpy.pi

        centralpixel = self.config.centralpixel  # (column, row) = (delta, gamma)
        beta_range = -app[1] * (numpy.arange(data.shape[1]) - centralpixel[1]) + beta
        delta_range = app[0] * (numpy.arange(data.shape[0]) - centralpixel[0]) + delta

        # masking
        if self.config.maskmatrix is not None:
            if self.config.maskmatrix.shape != data.shape:
                raise errors.BackendError(
                    "The mask matrix does not have the same shape as the images"
                )
            weights *= self.config.maskmatrix

        delta_range = delta_range[self.config.ymask]
        beta_range = beta_range[self.config.xmask]

        weights = self.apply_mask(weights, self.config.xmask, self.config.ymask)
        intensity = self.apply_mask(data, self.config.xmask, self.config.ymask)

        intensity = numpy.rot90(intensity)
        intensity = numpy.fliplr(intensity)
        intensity = numpy.flipud(intensity)

        weights = numpy.rot90(weights)
        weights = numpy.fliplr(weights)
        weights = numpy.flipud(weights)

        # polarisation correction
        delta_grid, beta_grid = numpy.meshgrid(delta_range, beta_range)
        Pver = (
            1
            - numpy.sin(delta_grid * numpy.pi / 180.0) ** 2
            * numpy.cos(beta_grid * numpy.pi / 180.0) ** 2
        )
        # intensity /= Pver

        return (
            intensity,
            weights,
            (wavelength, UB, beta_range, delta_range, omega, alfa, chi, phi),
        )

    def get_point_params(self, scan, first, last):
        sl = slice(first, last + 1)

        DEL, OME, ALF, BET, CHI, PHI, MON, TRANSM = list(range(8))
        params = numpy.zeros(
            (last - first + 1, 8)
        )  # gamma delta theta chi phi mu mon transm
        params[:, CHI] = 0  # scan.motorpos('CHI')
        params[:, PHI] = 0  # scan.motorpos('PHI')

        params[:, OME] = scan.datacol("omecnt")[sl]
        params[:, BET] = scan.datacol("betcnt")[sl]
        params[:, DEL] = scan.datacol("delcnt")[sl]
        params[:, MON] = scan.datacol("Monitor")[sl]

        # params[:, TRANSM] = scan.datacol('transm')[sl]
        params[:, TRANSM] = 1

        params[:, ALF] = scan.datacol("alfcnt")[sl]
        return params


def load_matrix(filename):
    if filename == None:
        return None
    if os.path.exists(filename):
        ext = os.path.splitext(filename)[-1]
        if ext == ".txt":
            return numpy.array(numpy.loadtxt(filename), dtype=numpy.bool)
        elif ext == ".npy":
            return numpy.array(numpy.load(filename), dtype=numpy.bool)
        elif ext == ".edf":
            return numpy.array(EdfFile.EdfFile(filename).getData(0), dtype=numpy.bool)
        else:
            raise ValueError(
                "unknown extension {0}, unable to load matrix!\n".format(ext)
            )
    else:
        raise IOError(
            "filename: {0} does not exist. Can not load matrix".format(filename)
        )
