# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
# All rights reserved.  Use of this source code is governed by
# a BSD-style license that can be found in the LICENSE file.

from .mpi import MPITestCase

import os

import numpy as np

from ..tod import OpPolyFilter, AnalyticNoise, OpSimNoise, Interval
from ..todmap import TODHpixSpiral

from ._helpers import create_outdir, create_distdata, boresight_focalplane


class OpPolyFilterTest(MPITestCase):
    def setUp(self):
        fixture_name = os.path.splitext(os.path.basename(__file__))[0]
        self.outdir = create_outdir(self.comm, fixture_name)

        # Create one observation per group, and each observation will have
        # one detector per process and a single chunk.  Data within an
        # observation is distributed by detector.

        self.data = create_distdata(self.comm, obs_per_group=1)
        self.ndet = self.data.comm.group_size
        self.rate = 20.0

        # Create detectors with a range of knee frequencies.
        (
            dnames,
            dquat,
            depsilon,
            drate,
            dnet,
            dfmin,
            dfknee,
            dalpha,
        ) = boresight_focalplane(
            self.ndet,
            samplerate=self.rate,
            net=10.0,
            fmin=1.0e-5,
            fknee=np.linspace(0.01, 0.19, num=self.ndet),
        )

        # Fitting order
        self.order = 5

        # Total Samples
        self.totsamp = 100000

        # Populate the observations (one per group)

        tod = TODHpixSpiral(
            self.data.comm.comm_group,
            dquat,
            self.totsamp,
            detranks=self.data.comm.group_size,
            firsttime=0.0,
            rate=self.rate,
            nside=512,
        )

        # Construct an analytic noise model for the detectors

        nse = AnalyticNoise(
            rate=drate,
            fmin=dfmin,
            detectors=dnames,
            fknee=dfknee,
            alpha=dalpha,
            NET=dnet,
        )

        intervals = []
        interval_len = 1000
        for istart in range(0, self.totsamp, interval_len):
            istop = min(istart + interval_len, self.totsamp)
            intervals.append(
                Interval(
                    start=istart / self.rate,
                    stop=istop / self.rate,
                    first=istart,
                    last=istop - 1,
                )
            )

        self.data.obs[0]["tod"] = tod
        self.data.obs[0]["noise"] = nse
        self.data.obs[0]["intervals"] = intervals

    def test_filter(self):
        # generate timestreams
        op = OpSimNoise()
        op.exec(self.data)

        # Replace the noise with a polynomial fit
        old_rms = []
        for ob in self.data.obs:
            tod = ob["tod"]
            orms = {}
            for det in tod.local_dets:
                cachename = "noise_{}".format(det)
                y = tod.cache.reference(cachename)
                x = np.arange(y.size)
                p = np.polyfit(x, y, self.order)
                y[:] = np.polyval(p, x)
                orms[det] = np.std(y)
                del y
            old_rms.append(orms)

        # Filter timestreams
        op = OpPolyFilter(name="noise", order=self.order)
        op.exec(self.data)

        # Ensure all timestreams are zeroed out by the filter.
        # The polynomial basis used in populating the TOD is different
        # than the one filter uses and the basis functions are not
        # strictly orthogonal on a sparse grid so we expect a low level
        # residual.

        for ob, orms in zip(self.data.obs, old_rms):
            tod = ob["tod"]
            for det in tod.local_dets:
                cachename = "noise_{}".format(det)
                y = tod.cache.reference(cachename)
                rms = np.std(y)
                old = orms[det]
                if np.abs(rms / old) > 1e-6:
                    raise RuntimeError(
                        "det {} old rms = {}, new rms = {}" "".format(det, old, rms)
                    )
                del y
        return
