#!python3
"""
Predict protein pKa based on MCCE method.
http://pka.engr.ccny.cuny.edu/

Require MCCE 3.0 to work: https://anaconda.org/SalahSalah/mcce/files
"""
import asyncio
import glob
import gzip
import locale
import logging
import math
import os
import re
import shutil
import subprocess
import sys
import time
from multiprocessing import Pool
from urllib.request import urlopen

import aioftp
import pandas as pd
import uvloop

# Sapelo Locale is broken, quick fix
locale.setlocale(locale.LC_ALL, "en_US.UTF-8")
# Set working directory
ROOTPATH = os.path.dirname(os.path.realpath(sys.argv[0]))
os.chdir(ROOTPATH)
# Log settings
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
handler = logging.FileHandler(f"./pKa_calculation_{__file__}.log")
handler.setLevel(logging.INFO)
formatter = logging.Formatter(
    "%(asctime)s\t%(levelname)s\t"
    "[%(filename)s:%(lineno)s -%(funcName)12s()]\t%(message)s"
)
handler.setFormatter(formatter)
logger.addHandler(handler)

class pdb:
    def __init__(self):
        self.all_ids = []
        self.download_ids = []  # Download -> Unzip -> Preprocess -> Calculate
        self.unzip_ids = []  # Unzip -> Preprocess -> Calculate
        self.preprocess_ids = []  # Preprocess -> Calculate
        self.ready_ids = []  # Calculate
        self.finished_ids = []  # Successfully calculated IDs
        self.error_ids = []  # Error in download, unzip, or calculation
        # IDs this script will work on (messy queue implementation)
        self.working_ids = []

    def load_id(self):
        """
        First try to get existing pKa values,
        then get the list of PDB files to download.
        """
        for folder in ["./pdb", "./annotation", "./results"]:
            try:
                os.makedirs(folder)
            except OSError:
                pass
        self.finished_ids = [id[-8:-4] for id in glob.glob("./results/*.pka")]
        logger.debug(f"{len(self.finished_ids)} finished files.")
        # Create file even at first run so that the results folder doesn't get deleted
        with open("./results/finished_ids.list", "a") as f:
            f.write("\n".join(self.finished_ids))

        self.ready_ids = list(set(
            [id[-12:-8].upper() for id in glob.glob("./pdb/*/*.pdb.bak")]) - set(self.finished_ids))
        logger.debug(f"{len(self.ready_ids)} files ready to be calculated.")

        self.preprocess_ids = list(set([id[-8:-4].upper() for id in glob.glob(
            "./pdb/*/*.pdb") if "out" not in id]) - set(self.finished_ids) - set(self.ready_ids))
        logger.debug(
            f"{len(self.preprocess_ids)} files ready to be preprocessed.")

        self.unzip_ids = [id[-11:-7].upper() for id in glob.glob("./*.ent.gz")]
        logger.debug(f"{len(self.unzip_ids)} files ready to be unzipped.")

        if not os.path.exists("./annotation/uniprot_id_mapping.dat"):
            with urlopen("ftp://ftp.uniprot.org/pub/databases/uniprot/current_release/knowledgebase/idmapping/by_organism/HUMAN_9606_idmapping.dat.gz") as remotefile:
                logger.debug(
                    "Saving UniProt ID mapping data since it doesn't exist...")
                with open("./annotation/uniprot_id_mapping.dat.gz", "wb") as f:
                    f.write(remotefile.read())
            with gzip.open(
                "./annotation/uniprot_id_mapping.dat.gz", "rb") as inFile, open(
                    "./annotation/uniprot_id_mapping.dat", "wb") as outFile:
                shutil.copyfileobj(inFile, outFile)
            os.remove("./annotation/uniprot_id_mapping.dat.gz")
        else:
            logger.debug("UniProt ID mapping data exists.")

        logger.debug("Reading all possible PDB IDs...")
        annot = pd.read_csv("./annotation/uniprot_id_mapping.dat",
                            sep="\t", header=None,
                            names=["uniprot", "id", "value"])
        self.all_ids = annot.loc[annot.id == "PDB", "value"].tolist()
        self.download_ids = list(set(self.all_ids) - set(self.unzip_ids) - set(
            self.preprocess_ids) - set(self.ready_ids) - set(self.finished_ids))
        logger.info(
            f"{len(self.download_ids)} PDB files need to be downloaded.")

    def get_link(self, ids):
        """ Get PDB file links from:
            ftp://ftp.wwpdb.org/pub/pdb/data/structures/divided/pdb/ ,
            and create folders to store the files.

        Parameters
        ----------
            ids: list
                The PDB IDs to download.

        Returns
        -------
            Links to download.
        """
        if isinstance(ids, list):
            ids = [id[:4].lower() for id in ids]  # pdb file IDs
            pdb_names = [f"{id}.ent.gz" for id in ids]  # pdb filenames
            # subdirectory of the pdb files
            pdbDirs = [id[1:3].lower() for id in ids]
            remoteaddr = [
                f"ftp://ftp.wwpdb.org/pub/pdb/data/structures/divided/pdb/{pdbDir}/pdb{pdb_name}" for pdbDir, pdb_name in zip(pdbDirs, pdb_names)]
        else:
            raise TypeError(f"{id} is not a string or list.")
        return remoteaddr

    def make_dirs(self, ids):
        """Make sure the download directory exists."""
        for id in ids:
            try:
                os.makedirs(os.path.join(ROOTPATH, "pdb", id.upper()))
            except OSError:
                pass

    async def download_worker(self, session, url):
        """Download the given url to working directory."""
        url = url[len("ftp://ftp.wwpdb.org"):]
        logger.debug(f"Downloading {url}")
        try:
            await session.download(url)
            self.unzip_ids.append(url[-11:-7].upper())
        except Exception as e:
            self.error_ids.append(url[-11:-7].upper())
            logger.warning(f"Error when downloading {url}: {e}")

    async def download_session(self, sem, work_queue):
        """ Get urls from the queue and pass to worker.

        Parameters
        ----------
            sem: asyncio.Semaphore object
            work_queue: asyncio.Queue object
        """
        while not work_queue.empty():
            url = await work_queue.get()
            logger.debug(f"Got url from queue: {url}")
            async with sem:
                async with aioftp.ClientSession("ftp.wwpdb.org") as session:
                    await self.download_worker(session, url)

    def download_queue(self, urls):
        """ Create a queue to download all the given urls.

        Parameters
        ----------
            urls: list
                A list of urls to download.

        Returns
        -------
            Downloaded "*.ent.gz" files in working directory.
        """
        logger.debug(f"{len(urls)} urls to download.")
        loop = uvloop.new_event_loop()
        asyncio.set_event_loop(loop)
        q = asyncio.Queue()
        sem = asyncio.Semaphore(10)
        [q.put_nowait(url) for url in urls]
        tasks = [asyncio.ensure_future(self.download_session(sem, q))
                 for _ in range(len(urls))]
        loop.run_until_complete(asyncio.gather(*tasks))
        # Zero-sleep to allow underlying connections to close
        loop.run_until_complete(asyncio.sleep(0))
        loop.close()

    def check_mcce(self):
        """Check if MCCE 3.0 exists."""
        if not os.path.exists(os.path.join(ROOTPATH, "mcce3.0")):
            if not os.path.exists(os.path.join(ROOTPATH, "mcce3.0.tar.bz2")):
                logger.debug("MCCE isn't downloaded yet. Retrieving...")
                with urlopen("https://anaconda.org/SalahSalah/mcce/3.0/download/linux-64/mcce-3.0-0.tar.bz2") as remotefile:
                    with open("./mcce-3.0-0.tar.bz2", 'wb') as f:
                        f.write(remotefile.read())
            subprocess.run(["tar", "-xjf", "mcce-3.0-0.tar.bz2"])
            shutil.move("./info/recipe/mcce3.0", "./mcce3.0")
            shutil.rmtree(os.path.join(ROOTPATH, "info"), ignore_errors=True)
            shutil.rmtree(os.path.join(ROOTPATH, "bin"), ignore_errors=True)
        else:
            logger.info("MCCE 3.0 exists, proceeding to calculation...")

    def unzip(self, id):
        """Unzip downloaded *.ent.gz file."""
        try:
            saved_pdb = os.path.join(ROOTPATH, "pdb", id, f"{id}.pdb")
            with gzip.open(f"pdb{id.lower()}.ent.gz", "rb") as inFile, open(saved_pdb, "wb") as outFile:
                shutil.copyfileobj(inFile, outFile)
            os.remove(f"pdb{id.lower()}.ent.gz")
            self.preprocess_ids.append(id)
        except Exception as e:
            self.error_ids.append(id)
            logger.warning(f"Unzip of {id} unsuccessful: {e}")

    def preprocess(self, id, backup=True):
        """
        This program will:
        1) strip lines other than ATOM and HETATM records
        2) keep the first model of an NMR structure
        3) delete H and D atoms
        4) MSE to MET residue
        5) keep only one atom alternate position
        6) keep defined chains, if chain ID(s) are given in command
        7) remove some cofactors and salt ions

        Parameters
        ----------
            id: str
                The PDB ID to find the file.
            backup: bool, optional
                Whether to backup the original file or not. Default is True,
                and save to "original.bak".

        Returns
        -------
            Nothing, modify the file in place.
        """
        removable_res = [
            " ZN", "PCA", "XYP", " NA", " CL", " CA", " MG", " MN", "HOH"
        ]
        model_start = False
        newlines = []
        ID = id.upper()
        filepath = os.path.join(ROOTPATH, "pdb", ID, f"{ID}.pdb")
        if backup:
            shutil.copy2(filepath, f"{filepath}.bak")
        with open(filepath) as f:
            for line in f:
                if line[:5] == "MODEL":
                    model_start = True
                if model_start and line[:6] == "ENDMDL":
                    break
                if line[:6] != "ATOM  " and line[:6] != "HETATM":
                    continue  # discard non ATOM records
                if line[13] == "H" or line[12] == "H":
                    continue
                if line[16] == "A":
                    line = f"{line[:16]} {line[17:]}"
                elif line[16] != " ":
                    continue  # delete this line, alternative posion is not A or empty
                if line[:6] == "HETATM" and line[17:20] == "MSE":
                    if line[12:15] == "SE ":
                        line = f"ATOM  {line[6:12]} SD{line[15:17]}MET{line[20:]}"
                    else:
                        line = f"ATOM  {line[6:17]}MET{line[20:]}"
                res = line[17:20]
                if res in removable_res:
                    continue
                newlines.append(line.rstrip())
        with open(filepath, "w") as f:
            f.write("\n".join(newlines))
        logger.debug(f"{ID} preprocessing complete.")

    def set_params(self, id, quickrun=True):
        """
        Set the parameters for MCCE.

        Parameters
        ----------
            id: str
                The PDB ID of the file.
            quickrun: bool, optional
                Use "run.prm.quick" or "run.prm.default".

        Returns
        -------
            run.prm: a file describing the parameters that points to the PDB file.
        """
        pkgpath = os.path.join(ROOTPATH, "mcce3.0")
        ID = id.upper()
        filepath = os.path.join(ROOTPATH, "pdb", ID)
        newlines = []
        if quickrun:
            shutil.copy2(
                os.path.join(pkgpath, "run.prm.quick"),
                os.path.join(filepath, "run.prm")
            )
        else:
            shutil.copy2([
                os.path.join(pkgpath, "run.prm.default"),
                os.path.join(filepath, "run.prm")
            ])
        with open(os.path.join(filepath, "run.prm")) as f:
            for line in f:
                line = line.rstrip()
                if line.endswith("(INPDB)"):
                    line = re.sub(r"^[^\s]+", fr"{id}.pdb", line)
                if line.endswith(("(DO_PREMCCE)", "(DO_ROTAMERS)",
                                  "(DO_ENERGY)", "(DO_MONTE)")):
                    line = re.sub(r"^f", r"t", line)
                if line.endswith("(EPSILON_PROT)"):
                    line = re.sub(r"^[\d\.]+", r"8.0", line)
                if line.startswith("/home/mcce/mcce3.0"):
                    line = re.sub(r"^/.*3\.0", pkgpath,
                                  line)
                newlines.append(line)
        with open(os.path.join(filepath, "run.prm"), "w") as f:
            f.write("\n".join(newlines))
        self.ready_ids.append(ID)
        logger.debug(f"Parameters set for {ID}.")

    def split_ready_ids(self, num):
        """ A naive queue implementation for multiple scripts.

        Parameters
        ----------
            num: int
                Which part of the IDs to work on.

        Returns
        -------
            A list of the actual IDs to work on, and save the lists of IDs for
            other scripts to work with if this is the first instance.
        """
        if os.path.isfile(os.path.join(ROOTPATH, "results", "working_ids.list")):
            with open(os.path.join(ROOTPATH, "results", f"working_ids.list{num}"), "r") as f:
                self.working_ids = [line.strip() for line in f]
        else:
            n = math.ceil(len(self.ready_ids) / 10)
            self.working_ids = [self.ready_ids[i:i + n]
                                for i in range(0, len(self.ready_ids), n)]
            metafile = []
            for i, ids in enumerate(self.working_ids):
                metafile.append(os.path.join(
                    ROOTPATH, "results", f"working_ids.list{i}"))
                with open(os.path.join(ROOTPATH, "results", f"working_ids.list{i}"), "w") as f:
                    f.write("\n".join(ids))
                logger.debug(
                    f"Saved {len(ids)} IDs to file working_ids.list{i} .")
            with open(os.path.join(ROOTPATH, "results", "working_ids.list"), "w") as f:
                f.write("\n".join(metafile))
            self.working_ids = self.working_ids[num]

    def calc_pka(self, id, clean=True):
        """ Calculate protein pKa values using MCCE.
            https://sites.google.com/site/mccewiki/home

        Parameters
        ----------
            id: str
                The PDB ID of the protein calculated.
            clean: bool, optional
                Only keep the PDB file, run log and pKa output.

        Returns
        -------
            A set of files in a subdirectory named after the ID.
            See user manual for detail.
        """
        id = id.upper()
        os.chdir(os.path.realpath(os.path.join(ROOTPATH, "pdb", id)))
        logger.info(f"{id} calculation started.")
        start = time.time()
        with open(f"{id}.run.log", "w") as f:
            subprocess.run(f"{ROOTPATH}/mcce3.0/mcce", stdout=f)
        with open(f"{id}.run.log", "rb") as f:
            last = f.readlines()[-1].decode().lstrip()
        if last.startswith(("Fatal", "FATAL", "WARNING", "STOP")):
            self.error_ids.append(id)
            logger.warning(
                f"{id} calculation aborted after {time.time() - start}s, due to {last}")
        else:
            self.finished_ids.append(id)
            logger.info(
                f"{id} calculation finished, used {time.time() - start}s.")
            shutil.move("pK.out", os.path.join(
                ROOTPATH, "results", f"{id}.pka"))
        if clean:
            del_list = [i for i in os.listdir() if i not in (
                "pK.out", f"{id}.run.log", f"{id}.pdb.bak")]
            [os.remove(item) for item in del_list]

if __name__ == "__main__":
    x = pdb()
    x.load_id()
    urls = x.get_link(x.download_ids)
    x.make_dirs(x.all_ids)
    x.download_queue(urls)

    x.check_mcce()
    for id in x.unzip_ids:
        x.unzip(id)
    for id in x.preprocess_ids:
        try:
            x.preprocess(id)
            x.set_params(id)
        except Exception as e:
            x.error_ids.append(id)
            logger.warning(f"Preprocess of {id}: {e}")
    # subprocess.run(["find", ".", "-type", "d", "-empty", "-delete"])

    x.split_ready_ids(0)  # 0 - 9, run 0 first to generate other lists

    with Pool(os.cpu_count()) as p:
        p.map(x.calc_pka, x.working_ids)

    with open("./results/finished_ids.list", "a") as f:
        f.write("\n".join(x.working_ids))

    with open("./results/error_ids.list", "a") as f:
        f.write("\n".join(x.error_ids))


#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt

import ld as LD
import dd as DD

from hidespines import *

import sys

''' compares difference between Sn and moment equations as cell width --> 0 '''

if (len(sys.argv) > 1):
	outfile = sys.argv[1]
else:
	outfile = None

def getDiff(sol, tol=1e-6):

	diff = np.zeros(len(sol))
	for i in range(len(sol)):

		x, phi, it = sol[i].sourceIteration(tol)

		# diff[i] = np.linalg.norm(phi - sol[i].phi_SN, 2)/np.linalg.norm(sol[i].phi_SN, 2)
		diff[i] = np.linalg.norm(phi - sol[i].phi_SN, 2)/np.linalg.norm(sol[i].phi_SN, 2)

	return diff

N = 100
n = 8
xb = 1

Sigmaa = lambda x: .1
Sigmat = lambda x: 1
q = lambda x, mu: 1

tol = 1e-10

N = np.logspace(1, 3, 5)

N = np.array([int(x) for x in N])

ed00 = [LD.Eddington(np.linspace(0, xb, x+1), n, Sigmaa,
	Sigmat, q, OPT=0, GAUSS=0) for x in N]

ed01 = [LD.Eddington(np.linspace(0, xb, x+1), n, Sigmaa,
	Sigmat, q, OPT=0, GAUSS=1) for x in N]

ed10 = [LD.Eddington(np.linspace(0, xb, x+1), n, Sigmaa,
	Sigmat, q, OPT=1, GAUSS=0) for x in N]

ed11 = [LD.Eddington(np.linspace(0, xb, x+1), n, Sigmaa,
	Sigmat, q, OPT=1, GAUSS=1) for x in N]

ed20 = [LD.Eddington(np.linspace(0, xb, x+1), n, Sigmaa,
	Sigmat, q, OPT=2, GAUSS=0) for x in N]

ed21 = [LD.Eddington(np.linspace(0, xb, x+1), n, Sigmaa,
	Sigmat, q, OPT=2, GAUSS=1) for x in N]

diff00 = getDiff(ed00, tol)
diff01 = getDiff(ed01, tol)
diff10 = getDiff(ed10, tol)
diff11 = getDiff(ed11, tol)
diff20 = getDiff(ed20, tol)
diff21 = getDiff(ed21, tol)

fontsize=16
plt.loglog(xb/N, diff00, '-o', clip_on=False, label='MHFEM Edges, No Gauss')
plt.loglog(xb/N, diff01, '-o', clip_on=False, label='Maintain Slopes, No Gauss')
plt.loglog(xb/N, diff10, '-o', clip_on=False, label='MHFEM Edges, Gauss')
plt.loglog(xb/N, diff11, '-o', clip_on=False, label='Maintain Slopes, Gauss')
plt.loglog(xb/N, diff20, '-o', clip_on=False, label='vanLeer, No Gauss')
plt.loglog(xb/N, diff21, '-o', clip_on=False, label='vanLeer, Gauss')
plt.xlabel(r'$h$', fontsize=fontsize)
plt.ylabel('SN/MHFEM Convergence', fontsize=fontsize)
plt.legend(loc='best', frameon=False)
hidespines(plt.gca())
if (outfile != None):
	plt.savefig(outfile, transparent=True)
else:
	plt.show()


from distutils.version import LooseVersion
from textwrap import dedent
from io import StringIO

import numpy
import scipy
from scipy import stats
import pandas

from unittest import mock
import pytest
import pandas.testing as pdtest
from wqio.tests import helpers

from wqio.features import Location, Dataset
from wqio.datacollections import DataCollection, _dist_compare

OLD_SCIPY = LooseVersion(scipy.version.version) < LooseVersion("0.19")

def check_stat(expected_csv, result, comp=False):
    index_col = [0]
    if comp:
        index_col += [1]

    file_obj = StringIO(dedent(expected_csv))
    expected = pandas.read_csv(file_obj, header=[0, 1], index_col=index_col)

    if comp:
        expected = expected.stack(level=-1)

    pdtest.assert_frame_equal(
        expected.sort_index(axis="columns"),
        result.sort_index(axis="columns").round(6),
        atol=1e-5,
    )

def remove_g_and_h(group):
    return group.name[1] not in ["G", "H"]

@pytest.fixture
def dc():
    df = helpers.make_dc_data_complex()
    dc = DataCollection(
        df,
        rescol="res",
        qualcol="qual",
        stationcol="loc",
        paramcol="param",
        ndval="<",
        othergroups=None,
        pairgroups=["state", "bmp"],
        useros=True,
        filterfxn=remove_g_and_h,
        bsiter=10000,
    )

    return dc

@pytest.fixture
def dc_noNDs():
    df = helpers.make_dc_data_complex()
    dc = DataCollection(
        df,
        rescol="res",
        qualcol="qual",
        stationcol="loc",
        paramcol="param",
        ndval="junk",
        othergroups=None,
        pairgroups=["state", "bmp"],
        useros=True,
        filterfxn=remove_g_and_h,
        bsiter=10000,
    )

    return dc

def test_basic_attr(dc):
    assert dc._raw_rescol == "res"
    assert isinstance(dc.data, pandas.DataFrame)
    assert dc.roscol == "ros_res"
    assert dc.rescol == "ros_res"
    assert dc.qualcol == "qual"
    assert dc.stationcol == "loc"
    assert dc.paramcol == "param"
    assert dc.ndval == ["<"]
    assert dc.bsiter == 10000
    assert dc.groupcols == ["loc", "param"]
    assert dc.tidy_columns == ["loc", "param", "res", "__censorship"]
    assert hasattr(dc, "filterfxn")

def test_data(dc):
    assert isinstance(dc.data, pandas.DataFrame)
    assert dc.data.shape == (519, 8)
    assert "G" in dc.data["param"].unique()
    assert "H" in dc.data["param"].unique()

@pytest.mark.parametrize("useros", [True, False])
def test_tidy(dc, useros):
    assert isinstance(dc.tidy, pandas.DataFrame)
    assert dc.tidy.shape == (388, 5)
    assert "G" not in dc.tidy["param"].unique()
    assert "H" not in dc.tidy["param"].unique()
    collist = ["loc", "param", "res", "__censorship", "ros_res"]
    assert dc.tidy.columns.tolist() == collist

def test_paired(dc):
    assert isinstance(dc.paired, pandas.DataFrame)
    assert dc.paired.shape == (164, 6)
    assert "G" not in dc.paired.index.get_level_values("param").unique()
    assert "H" not in dc.paired.index.get_level_values("param").unique()
    dc.paired.columns.tolist() == [
        ("res", "Inflow"),
        ("res", "Outflow"),
        ("res", "Reference"),
        ("__censorship", "Inflow"),
        ("__censorship", "Outflow"),
        ("__censorship", "Reference"),
    ]

def test_count(dc):
    known_csv = """\
        station,Inflow,Outflow,Reference
        result,Count,Count,Count
        param,,,
        A,21,22,20
        B,24,22,19
        C,24,24,25
        D,24,25,21
        E,19,16,20
        F,21,24,17
    """
    check_stat(known_csv, dc.count)

def test_n_unique(dc):
    known_csv = """\
        loc,Inflow,Outflow,Reference
        result,bmp,bmp,bmp
        param,,,
        A,7,7,7
        B,7,7,7
        C,7,7,7
        D,7,7,7
        E,7,7,7
        F,7,7,7
        G,7,7,7
        H,7,7,7
    """
    check_stat(known_csv, dc.n_unique("bmp"))

@helpers.seed
def test_median(dc):
    known_csv = """\
        station,Inflow,Inflow,Inflow,Outflow,Outflow,Outflow,Reference,Reference,Reference
        result,lower,median,upper,lower,median,upper,lower,median,upper
        param,,,,,,,,,
        A,0.334506,1.197251,2.013994,0.860493,2.231058,2.626023,1.073386,1.639472,1.717293
        B,1.366948,2.773989,3.297147,0.23201,1.546499,2.579206,0.204164,1.565076,2.196367
        C,0.17351,0.525957,0.68024,0.247769,0.396984,0.540742,0.136462,0.412693,0.559458
        D,0.374122,1.201892,2.098846,0.516989,1.362759,1.827087,0.314655,0.882695,1.24545
        E,0.276095,1.070858,1.152887,0.287914,0.516746,1.456859,0.366824,0.80716,2.040739
        F,0.05667,0.832488,1.310575,0.425237,1.510942,2.193997,0.162327,0.745993,1.992513
    """
    check_stat(known_csv, dc.median)

@helpers.seed
def test_mean(dc):
    known_csv = """\
        station,Inflow,Inflow,Inflow,Outflow,Outflow,Outflow,Reference,Reference,Reference
        result,lower,mean,upper,lower,mean,upper,lower,mean,upper
        param,,,,,,,,,
        A,1.231607,2.646682,4.204054,1.930601,5.249281,9.081952,1.540167,3.777974,6.389439
        B,2.99031,7.647175,12.810844,1.545539,6.863835,12.705913,1.010374,4.504255,9.592572
        C,0.37496,0.513248,0.65948,0.411501,1.004637,1.706317,0.35779,0.541962,0.734751
        D,1.29141,3.021235,4.987855,1.285899,2.318808,3.451824,1.008364,1.945828,2.924812
        E,0.818641,1.914696,3.049554,0.584826,1.098241,1.640807,1.113589,2.283292,3.581946
        F,0.8379,9.825404,25.289933,1.497825,3.450184,5.61929,0.939917,2.491708,4.094258
    """
    check_stat(known_csv, dc.mean)

@helpers.seed
def test_std_dev(dc):
    known_csv = """\
        station,Inflow,Outflow,Reference
        result,std. dev.,std. dev.,std. dev.
        param,,,
        A,3.58649,8.719371,5.527633
        B,12.360099,13.60243,10.759285
        C,0.353755,1.691208,0.493325
        D,4.811938,2.849393,2.248178
        E,2.55038,1.096698,2.789238
        F,34.447565,5.361033,3.398367
    """
    check_stat(known_csv, dc.std_dev)

@helpers.seed
def test_percentile_25(dc):
    known_csv = """\
        station,Inflow,Outflow,Reference
        result,pctl 25,pctl 25,pctl 25
        param,,,
        A,0.522601,0.906029,1.094721
        B,1.472541,0.251126,0.314226
        C,0.164015,0.267521,0.136462
        D,0.35688,0.516989,0.383895
        E,0.364748,0.311508,0.394658
        F,0.120068,0.406132,0.224429
    """
    check_stat(known_csv, dc.percentile(25))

@helpers.seed
def test_percentile_75(dc):
    known_csv = """\
        station,Inflow,Outflow,Reference
        result,pctl 75,pctl 75,pctl 75
        param,,,
        A,2.563541,3.838021,2.650648
        B,4.728871,2.849948,2.261847
        C,0.776388,0.853535,0.792612
        D,3.04268,2.79341,3.611793
        E,1.532775,1.59183,3.201534
        F,1.792985,2.80979,2.742249
    """
    check_stat(known_csv, dc.percentile(75))

@helpers.seed
def test_logmean(dc):
    known_csv = """\
        station,Inflow,Inflow,Inflow,Outflow,Outflow,Outflow,Reference,Reference,Reference
        result,Log-mean,lower,upper,Log-mean,lower,upper,Log-mean,lower,upper
        param,,,,,,,,,
        A,0.140559,-0.55112,0.644202,0.733004,0.047053,1.22099,0.545205,-0.057683,1.029948
        B,1.026473,0.368659,1.541241,0.105106,-0.939789,0.860244,0.068638,-0.932357,0.661203
        C,-0.963004,-1.304115,-0.638446,-0.83221,-1.464092,-0.414379,-1.088377,-1.556795,-0.720706
        D,0.062317,-0.663241,0.58349,0.185757,-0.325074,0.598432,-0.063507,-0.670456,0.434214
        E,-0.103655,-0.751075,0.385909,-0.456202,-1.08692,0.029967,-0.068135,-0.787007,0.51226
        F,-0.442721,-1.874677,0.344704,0.211658,-0.504166,0.734283,-0.253352,-1.175917,0.467231
    """
    check_stat(known_csv, dc.logmean)

@helpers.seed
def test_logstd_dev(dc):
    known_csv = """\
        station,Inflow,Outflow,Reference
        result,Log-std. dev.,Log-std. dev.,Log-std. dev.
        param,,,
        A,1.374026,1.343662,1.225352
        B,1.430381,2.07646,1.662001
        C,0.818504,1.263631,1.057177
        D,1.530871,1.187246,1.277927
        E,1.264403,1.121038,1.474431
        F,2.324063,1.516331,1.701596
    """
    check_stat(known_csv, dc.logstd_dev)

@helpers.seed
def test_geomean(dc):
    known_csv = """\
        station,Inflow,Inflow,Inflow,Outflow,Outflow,Outflow,Reference,Reference,Reference
        Geo-mean,Log-mean,lower,upper,Log-mean,lower,upper,Log-mean,lower,upper
        param,,,,,,,,,
        A,1.150917,0.576304,1.904467,2.081323,1.048178,3.390543,1.724962,0.943949,2.800919
        B,2.791205,1.445795,4.670381,1.110829,0.39071,2.363737,1.071049,0.393625,1.937121
        C,0.381744,0.271413,0.528113,0.435087,0.231288,0.66075,0.336763,0.210811,0.486409
        D,1.064299,0.515179,1.792283,1.204129,0.722474,1.819264,0.938467,0.511475,1.543749
        E,0.901536,0.471859,1.470951,0.633686,0.337254,1.03042,0.934134,0.455205,1.66906
        F,0.642286,0.153405,1.411572,1.235726,0.604009,2.083988,0.776195,0.308536,1.595571
    """
    check_stat(known_csv, dc.geomean)

@helpers.seed
def test_geostd_dev(dc):
    known_csv = """\
        station,Inflow,Outflow,Reference
        Geo-std. dev.,Log-std. dev.,Log-std. dev.,Log-std. dev.
        param,,,
        A,3.951225,3.833055,3.405365
        B,4.180294,7.976181,5.269843
        C,2.267105,3.538244,2.878234
        D,4.622199,3.278041,3.589191
        E,3.540977,3.068036,4.368548
        F,10.217099,4.55548,5.48269
    """
    check_stat(known_csv, dc.geostd_dev)

@helpers.seed
def test_shapiro(dc):
    known_csv = """\
        station,Inflow,Inflow,Outflow,Outflow,Reference,Reference
        result,pvalue,statistic,pvalue,statistic,pvalue,statistic
        param,,,,,,
        A,1.8e-05,0.685783,1e-06,0.576069,4e-06,0.61735
        B,1e-06,0.594411,0.0,0.530962,0.0,0.41471
        C,0.028774,0.905906,0.0,0.546626,0.00279,0.860373
        D,1e-06,0.622915,1.5e-05,0.722374,0.000202,0.76518
        E,1.7e-05,0.654137,0.004896,0.818813,0.000165,0.74917
        F,0.0,0.292916,2e-06,0.634671,0.000167,0.713968
    """
    check_stat(known_csv, dc.shapiro)

@helpers.seed
def test_shapiro_log(dc):
    known_csv = """\
        station,Inflow,Inflow,Outflow,Outflow,Reference,Reference
        result,statistic,pvalue,statistic,pvalue,statistic,pvalue
        param,,,,,,
        A,0.983521938,0.96662426,0.979861856,0.913820148,0.939460814,0.234214202
        B,0.957531095,0.390856266,0.97048676,0.722278714,0.967978418,0.735424638
        C,0.906479359,0.029602444,0.974698305,0.78197974,0.967106879,0.572929323
        D,0.989704251,0.995502174,0.990663111,0.997093379,0.964812279,0.617747009
        E,0.955088913,0.479993254,0.95211035,0.523841977,0.963425279,0.61430341
        F,0.97542423,0.847370088,0.982230783,0.933124721,0.966197193,0.749036908
    """
    check_stat(known_csv, dc.shapiro_log)

@helpers.seed
def test_lilliefors(dc):
    known_csv = """\
        station,Inflow,Inflow,Outflow,Outflow,Reference,Reference
        result,lilliefors,pvalue,lilliefors,pvalue,lilliefors,pvalue
        param,,,,,,
        A,0.308131,1.4e-05,0.340594,0.0,0.364453,0.0
        B,0.36764,0.0,0.420343,0.0,0.417165,0.0
        C,0.166799,0.082737,0.324733,0.0,0.161753,0.090455
        D,0.273012,6.7e-05,0.240311,0.000665,0.296919,3.7e-05
        E,0.341398,3e-06,0.239314,0.014862,0.233773,0.005474
        F,0.419545,0.0,0.331315,0.0,0.284249,0.000741
    """
    check_stat(known_csv, dc.lilliefors)

@helpers.seed
def test_lilliefors_log(dc):
    known_csv = """\
        station,Inflow,Inflow,Outflow,Outflow,Reference,Reference
        result,log-lilliefors,pvalue,log-lilliefors,pvalue,log-lilliefors,pvalue
        param,,,,,,
        A,0.08548109,0.95458004,0.15443943,0.19715747,0.20141389,0.03268737
        B,0.16162839,0.10505016,0.12447902,0.49697902,0.15934334,0.22969362
        C,0.16957278,0.07248915,0.12388174,0.44379732,0.11746642,0.48915671
        D,0.06885549,0.99,0.06067356,0.99,0.13401954,0.41967483
        E,0.13506577,0.47186822,0.14552341,0.47797919,0.09164876,0.92860794
        F,0.14420794,0.30694533,0.08463267,0.92741885,0.08586933,0.9800294
    """
    check_stat(known_csv, dc.lilliefors_log)

@helpers.seed
def test_anderson_darling(dc):
    with helpers.raises(NotImplementedError):
        _ = dc.anderson_darling

@helpers.seed
def test_anderson_darling_log(dc):
    with helpers.raises(NotImplementedError):
        _ = dc.anderson_darling_log

@helpers.seed
def test_mann_whitney(dc):
    known_csv = """\
        ,,mann_whitney,mann_whitney,mann_whitney,pvalue,pvalue,pvalue
        loc_2,,Inflow,Outflow,Reference,Inflow,Outflow,Reference
        param,loc_1,,,,,,
        A,Inflow,,180.0,179.0,,0.2198330905,0.4263216587
        A,Outflow,282.0,,248.0,0.2198330905,,0.488580368
        A,Reference,241.0,192.0,,0.4263216587,0.488580368,
        B,Inflow,,345.0,317.0,,0.0766949991,0.0304383994
        B,Outflow,183.0,,216.0,0.0766949991,,0.8650586835
        B,Reference,139.0,202.0,,0.0304383994,0.8650586835,
        C,Inflow,,282.0,323.0,,0.9097070273,0.6527104406
        C,Outflow,294.0,,323.0,0.9097070273,,0.6527104406
        C,Reference,277.0,277.0,,0.6527104406,0.6527104406,
        D,Inflow,,285.0,263.0,,0.7718162376,0.8111960975
        D,Outflow,315.0,,293.0,0.7718162376,,0.5082395211
        D,Reference,241.0,232.0,,0.8111960975,0.5082395211,
        E,Inflow,,164.0,188.0,,0.7033493939,0.9663820218
        E,Outflow,140.0,,132.0,0.7033493939,,0.3813114322
        E,Reference,192.0,188.0,,0.9663820218,0.3813114322,
        F,Inflow,,201.0,172.0,,0.2505911218,0.8601783903
        F,Outflow,303.0,,236.0,0.2505911218,,0.4045186043
        F,Reference,185.0,172.0,,0.8601783903,0.4045186043
    """
    check_stat(known_csv, dc.mann_whitney, comp=True)

@helpers.seed
def test_t_test(dc):
    known_csv = """\
        ,,pvalue,pvalue,pvalue,t_test,t_test,t_test
        loc_2,,Inflow,Outflow,Reference,Inflow,Outflow,Reference
        param,loc_1,,,,,,
        A,Inflow,,0.2178424157,0.4563196599,,-1.2604458127,-0.7539785777
        A,Outflow,0.2178424157,,0.5240147979,1.2604458127,,0.643450194
        A,Reference,0.4563196599,0.5240147979,,0.7539785777,-0.643450194,
        B,Inflow,,0.8430007638,0.3898358794,,0.1992705833,0.869235357
        B,Outflow,0.8430007638,,0.5491097882,-0.1992705833,,0.6043850808
        B,Reference,0.3898358794,0.5491097882,,-0.869235357,-0.6043850808,
        C,Inflow,,0.1847386316,0.8191392537,,-1.3639360123,-0.2300373632
        C,Outflow,0.1847386316,,0.2179907667,1.3639360123,,1.2615982727
        C,Reference,0.8191392537,0.2179907667,,0.2300373632,-1.2615982727,
        D,Inflow,,0.5484265023,0.344783812,,0.6056706932,0.9582600001
        D,Outflow,0.5484265023,,0.6299742693,-0.6056706932,,0.4851636024
        D,Reference,0.344783812,0.6299742693,,-0.9582600001,-0.4851636024,
        E,Inflow,,0.2304569921,0.6770414622,,1.2287029977,-0.4198288251
        E,Outflow,0.2304569921,,0.1023435465,-1.2287029977,,-1.6935358498
        E,Reference,0.6770414622,0.1023435465,,0.4198288251,1.6935358498,
        F,Inflow,,0.422008391,0.3549979666,,0.8190789273,0.9463539528
        F,Outflow,0.422008391,,0.4988994144,-0.8190789273,,0.6826435968
        F,Reference,0.3549979666,0.4988994144,,-0.9463539528,-0.6826435968
    """
    check_stat(known_csv, dc.t_test, comp=True)

@helpers.seed
def test_levene(dc):
    known_csv = """\
        ,,levene,levene,levene,pvalue,pvalue,pvalue
        loc_2,,Inflow,Outflow,Reference,Inflow,Outflow,Reference
        param,loc_1,,,,,,
        A,Inflow,,1.176282059,0.293152155,,0.284450688,0.591287419
        A,Outflow,1.176282059,,0.397705309,0.284450688,,0.531863542
        A,Reference,0.293152155,0.397705309,,0.591287419,0.531863542,
        B,Inflow,,0.003559637,0.402002411,,0.952694449,0.529578712
        B,Outflow,0.003559637,,0.408938588,0.952694449,,0.526247443
        B,Reference,0.402002411,0.408938588,,0.529578712,0.526247443,
        C,Inflow,,1.965613561,0.679535532,,0.167626459,0.413910674
        C,Outflow,1.965613561,,1.462364363,0.167626459,,0.232602352
        C,Reference,0.679535532,1.462364363,,0.413910674,0.232602352,
        D,Inflow,,0.643364813,0.983777911,,0.426532092,0.32681669
        D,Outflow,0.643364813,,0.116830634,0.426532092,,0.734124856
        D,Reference,0.983777911,0.116830634,,0.32681669,0.734124856,
        E,Inflow,,0.961616536,0.410491665,,0.333914902,0.525668596
        E,Outflow,0.961616536,,2.726351564,0.333914902,,0.107912818
        E,Reference,0.410491665,2.726351564,,0.525668596,0.107912818,
        F,Inflow,,0.841984453,0.734809611,,0.363948105,0.396999375
        F,Outflow,0.841984453,,0.25881357,0.363948105,,0.613802541
        F,Reference,0.734809611,0.25881357,,0.396999375,0.613802541,
    """
    check_stat(known_csv, dc.levene, comp=True)

@helpers.seed
def test_wilcoxon(dc):
    known_csv = """\
        ,,wilcoxon,wilcoxon,wilcoxon,pvalue,pvalue,pvalue
        loc_2,,Inflow,Outflow,Reference,Inflow,Outflow,Reference
        param,loc_1,,,,,,
        A,Inflow,,32.0,59.0,,0.03479,0.430679
        A,Outflow,32.0,,46.0,0.03479,,0.274445
        A,Reference,59.0,46.0,,0.430679,0.274445,
        B,Inflow,,38.0,22.0,,0.600179,0.182338
        B,Outflow,38.0,,31.0,0.600179,,0.858863
        B,Reference,22.0,31.0,,0.182338,0.858863,
        C,Inflow,,75.0,120.0,,0.167807,0.601046
        C,Outflow,75.0,,113.0,0.167807,,0.463381
        C,Reference,120.0,113.0,,0.601046,0.463381,
        D,Inflow,,44.0,31.0,,0.593618,0.530285
        D,Outflow,44.0,,45.0,0.593618,,0.972125
        D,Reference,31.0,45.0,,0.530285,0.972125,
        E,Inflow,,21.0,19.0,,0.910156,0.386271
        E,Outflow,21.0,,16.0,0.910156,,0.077148
        E,Reference,19.0,16.0,,0.386271,0.077148,
        F,Inflow,,62.0,22.0,,0.492459,0.952765
        F,Outflow,62.0,,28.0,0.492459,,0.656642
        F,Reference,22.0,28.0,,0.952765,0.656642,
    """
    with pytest.warns(UserWarning):
        check_stat(known_csv, dc.wilcoxon, comp=True)

@helpers.seed
def test_ranksums(dc):
    known_csv = """\
        ,,pvalue,pvalue,pvalue,rank_sums,rank_sums,rank_sums
        loc_2,,Inflow,Outflow,Reference,Inflow,Outflow,Reference
        param,loc_1,,,,,,
        A,Inflow,,0.2153009,0.4187782,,-1.2391203,-0.8085428
        A,Outflow,0.2153009,,0.4807102,1.2391203,,0.7051607
        A,Reference,0.4187782,0.4807102,,0.8085428,-0.7051607,
        B,Inflow,,0.0748817,0.029513,,1.781188,2.1765661
        B,Outflow,0.0748817,,0.8547898,-1.781188,,0.1830104
        B,Reference,0.029513,0.8547898,,-2.1765661,-0.1830104,
        C,Inflow,,0.9015386,0.6455162,,-0.1237179,0.46
        C,Outflow,0.9015386,,0.6455162,0.1237179,,0.46
        C,Reference,0.6455162,0.6455162,,-0.46,-0.46,
        D,Inflow,,0.7641772,0.8023873,,-0.3,0.2502587
        D,Outflow,0.7641772,,0.5011969,0.3,,0.6726078
        D,Reference,0.8023873,0.5011969,,-0.2502587,-0.6726078,
        E,Inflow,,0.6911022,0.9551863,,0.3973597,-0.0561951
        E,Outflow,0.6911022,,0.3727144,-0.3973597,,-0.8914004
        E,Reference,0.9551863,0.3727144,,0.0561951,0.8914004,
        F,Inflow,,0.2459307,0.8486619,,-1.1602902,-0.190826
        F,Outflow,0.2459307,,0.3971011,1.1602902,,0.8468098
        F,Reference,0.8486619,0.3971011,,0.190826,-0.8468098,
    """
    check_stat(known_csv, dc.ranksums, comp=True)

@helpers.seed
@pytest.mark.xfail(OLD_SCIPY, reason="Scipy < 0.19")
def test_kendall(dc):
    known_csv = """\
        ,,kendalltau,kendalltau,kendalltau,pvalue,pvalue,pvalue
        loc_2,,Inflow,Outflow,Reference,Inflow,Outflow,Reference
        param,loc_1,,,,,,
        A,Inflow,,-0.051661,-0.00738,,0.772893,0.967114
        A,Outflow,-0.051661,,-0.083333,0.772893,,0.690095
        A,Reference,-0.00738,-0.083333,,0.967114,0.690095,
        B,Inflow,,0.441351,0.298246,,0.015267,0.119265
        B,Outflow,0.441351,,0.559855,0.015267,,0.004202
        B,Reference,0.298246,0.559855,,0.119265,0.004202,
        C,Inflow,,0.280223,0.084006,,0.078682,0.578003
        C,Outflow,0.280223,,-0.1417,0.078682,,0.352394
        C,Reference,0.084006,-0.1417,,0.578003,0.352394,
        D,Inflow,,0.403469,0.095299,,0.020143,0.634826
        D,Outflow,0.403469,,0.318337,0.020143,,0.094723
        D,Reference,0.095299,0.318337,,0.634826,0.094723,
        E,Inflow,,0.114286,0.640703,,0.673337,0.004476
        E,Outflow,0.114286,,0.167944,0.673337,,0.449603
        E,Reference,0.640703,0.167944,,0.004476,0.449603,
        F,Inflow,,0.0,0.07231,,1.0,0.763851
        F,Outflow,0.0,,0.388889,1.0,,0.063
        F,Reference,0.07231,0.388889,,0.763851,0.063,
    """
    check_stat(known_csv, dc.kendall, comp=True)

@helpers.seed
def test_spearman(dc):
    known_csv = """\
        ,,pvalue,pvalue,pvalue,spearmanrho,spearmanrho,spearmanrho
        loc_2,,Inflow,Outflow,Reference,Inflow,Outflow,Reference
        param,loc_1,,,,,,
        A,Inflow,,0.7574884491,0.9627447553,,-0.0809319588,0.012262418
        A,Outflow,0.7574884491,,0.7617330788,-0.0809319588,,-0.0823529412
        A,Reference,0.9627447553,0.7617330788,,0.012262418,-0.0823529412,
        B,Inflow,,0.0110829791,0.0775159774,,0.5831305575,0.4537313433
        B,Outflow,0.0110829791,,0.0024069317,0.5831305575,,0.6850916941
        B,Reference,0.0775159774,0.0024069317,,0.4537313433,0.6850916941,
        C,Inflow,,0.1330504059,0.6063501968,,0.3387640122,0.1134228342
        C,Outflow,0.1330504059,,0.3431640379,0.3387640122,,-0.2070506455
        C,Reference,0.6063501968,0.3431640379,,0.1134228342,-0.2070506455,
        D,Inflow,,0.0195715066,0.4751861062,,0.4935814032,0.1858231711
        D,Outflow,0.0195715066,,0.1263974782,0.4935814032,,0.363209462
        D,Reference,0.4751861062,0.1263974782,,0.1858231711,0.363209462,
        E,Inflow,,0.9828818202,0.0013596162,,0.0084033613,0.8112988341
        E,Outflow,0.9828818202,,0.3413722947,0.0084033613,,0.3012263814
        E,Reference,0.0013596162,0.3413722947,,0.8112988341,0.3012263814,
        F,Inflow,,0.9645303744,0.6759971848,,-0.0106277141,0.1348767061
        F,Outflow,0.9645303744,,0.0560590794,-0.0106277141,,0.5028571429
        F,Reference,0.6759971848,0.0560590794,,0.1348767061,0.5028571429
    """
    check_stat(known_csv, dc.spearman, comp=True)

@helpers.seed
def test_theilslopes(dc):
    with helpers.raises(NotImplementedError):
        _ = dc.theilslopes

def test_inventory(dc):
    known_csv = StringIO(
        dedent(
            """\
    loc,param,Count,Non-Detect
    Inflow,A,21,3
    Inflow,B,24,6
    Inflow,C,24,0
    Inflow,D,24,11
    Inflow,E,19,4
    Inflow,F,21,8
    Outflow,A,22,1
    Outflow,B,22,9
    Outflow,C,24,4
    Outflow,D,25,12
    Outflow,E,16,2
    Outflow,F,24,8
    Reference,A,20,2
    Reference,B,19,6
    Reference,C,25,4
    Reference,D,21,12
    Reference,E,20,3
    Reference,F,17,7
    """
        )
    )
    expected = pandas.read_csv(known_csv, index_col=[0, 1]).astype(int)
    pdtest.assert_frame_equal(expected, dc.inventory.astype(int), check_names=False)

def test_inventory_noNDs(dc_noNDs):
    known_csv = StringIO(
        dedent(
            """\
    loc,param,Count,Non-Detect
    Inflow,A,21,0
    Inflow,B,24,0
    Inflow,C,24,0
    Inflow,D,24,0
    Inflow,E,19,0
    Inflow,F,21,0
    Outflow,A,22,0
    Outflow,B,22,0
    Outflow,C,24,0
    Outflow,D,25,0
    Outflow,E,16,0
    Outflow,F,24,0
    Reference,A,20,0
    Reference,B,19,0
    Reference,C,25,0
    Reference,D,21,0
    Reference,E,20,0
    Reference,F,17,0
    """
        )
    )
    expected = pandas.read_csv(known_csv, index_col=[0, 1]).astype(int)
    pdtest.assert_frame_equal(
        expected, dc_noNDs.inventory.astype(int), check_names=False,
    )

@helpers.seed
def test_stat_summary(dc):
    known_csv = StringIO(
        dedent(
            """\
    ros_res,loc,A,B,C,D,E,F
    Count,Inflow,21,24,24,24,19,21
    Count,Outflow,22,22,24,25,16,24
    Count,Reference,20,19,25,21,20,17
    Non-Detect,Inflow,3.0,6.0,0.0,11.0,4.0,8.0
    Non-Detect,Outflow,1.0,9.0,4.0,12.0,2.0,8.0
    Non-Detect,Reference,2.0,6.0,4.0,12.0,3.0,7.0
    mean,Inflow,2.64668,7.64717,0.51325,3.02124,1.9147,9.8254
    mean,Outflow,5.24928,6.86384,1.00464,2.31881,1.09824,3.45018
    mean,Reference,3.77797,4.50425,0.54196,1.94583,2.28329,2.49171
    std,Inflow,3.67506,12.62594,0.36136,4.91543,2.62027,35.29825
    std,Outflow,8.92456,13.92253,1.72758,2.90815,1.13267,5.47634
    std,Reference,5.67123,11.05411,0.5035,2.3037,2.8617,3.50296
    min,Inflow,0.0756,0.17404,0.10213,0.05365,0.08312,0.00803
    min,Outflow,0.11177,0.02106,0.03578,0.11678,0.07425,0.06377
    min,Reference,0.15575,0.04909,0.04046,0.08437,0.05237,0.03445
    10%,Inflow,0.1772,0.45233,0.13467,0.15495,0.1763,0.03548
    10%,Outflow,0.44852,0.08297,0.08222,0.26949,0.19903,0.18008
    10%,Reference,0.38448,0.13467,0.08241,0.19355,0.12777,0.09457
    25%,Inflow,0.5226,1.47254,0.16401,0.35688,0.36475,0.12007
    25%,Outflow,0.90603,0.25113,0.26752,0.51699,0.31151,0.40613
    25%,Reference,1.09472,0.31423,0.13646,0.3839,0.39466,0.22443
    50%,Inflow,1.19725,2.77399,0.52596,1.20189,1.07086,0.83249
    50%,Outflow,2.23106,1.5465,0.39698,1.36276,0.51675,1.51094
    50%,Reference,1.63947,1.56508,0.41269,0.8827,0.80716,0.74599
    75%,Inflow,2.56354,4.72887,0.77639,3.04268,1.53278,1.79299
    75%,Outflow,3.83802,2.84995,0.85354,2.79341,1.59183,2.80979
    75%,Reference,2.65065,2.26185,0.79261,3.61179,3.20153,2.74225
    90%,Inflow,6.02835,24.40655,0.99293,8.00691,6.28345,8.51706
    90%,Outflow,12.43052,23.90022,2.43829,5.66731,2.30348,10.32829
    90%,Reference,12.58278,6.67125,1.2205,4.78255,7.72012,8.57303
    max,Inflow,13.87664,45.97893,1.26657,21.75505,8.88365,163.01001
    max,Outflow,36.58941,47.49381,8.04948,12.39894,4.19118,23.29367
    max,Reference,21.22363,48.23615,1.94442,7.67751,8.75609,10.5095
    """
        )
    )

    expected = pandas.read_csv(known_csv, index_col=[0, 1]).T
    pdtest.assert_frame_equal(
        expected.round(5),
        dc.stat_summary().round(5),
        check_names=False,
        check_dtype=False,
        rtol=1e-4,
    )

def test_locations(dc):
    for loc in dc.locations:
        assert isinstance(loc, Location)
    assert len(dc.locations) == 18
    assert dc.locations[0].definition == {"loc": "Inflow", "param": "A"}
    assert dc.locations[1].definition == {"loc": "Inflow", "param": "B"}

def test_datasets(dc):
    _ds = []
    for d in dc.datasets("Inflow", "Outflow"):
        assert isinstance(d, Dataset)
        _ds.append(d)
    assert len(_ds) == 6
    assert _ds[0].definition == {"param": "A"}
    assert _ds[1].definition == {"param": "B"}

# this sufficiently tests dc._filter_collection
def test_selectLocations(dc):
    locs = dc.selectLocations(param="A", loc=["Inflow", "Outflow"])
    assert len(locs) == 2
    for n, (loc, loctype) in enumerate(zip(locs, ["Inflow", "Outflow"])):
        assert isinstance(loc, Location)
        assert loc.definition["param"] == "A"
        assert loc.definition["loc"] == loctype

def test_selectLocations_squeeze_False(dc):
    locs = dc.selectLocations(param="A", loc=["Inflow"], squeeze=False)
    assert len(locs) == 1
    for n, loc in enumerate(locs):
        assert isinstance(loc, Location)
        assert loc.definition["param"] == "A"
        assert loc.definition["loc"] == "Inflow"

def test_selectLocations_squeeze_True(dc):
    loc = dc.selectLocations(param="A", loc=["Inflow"], squeeze=True)
    assert isinstance(loc, Location)
    assert loc.definition["param"] == "A"
    assert loc.definition["loc"] == "Inflow"

def test_selectLocations_squeeze_True_None(dc):
    loc = dc.selectLocations(param="A", loc=["Junk"], squeeze=True)
    assert loc is None

# since the test_selectLocations* tests stress _filter_collection
# enough, we'll mock it out for datasets:
def test_selectDatasets(dc):
    with mock.patch.object(dc, "_filter_collection") as _fc:
        with mock.patch.object(dc, "datasets", return_value=["A", "B"]) as _ds:
            dc.selectDatasets("Inflow", "Reference", foo="A", bar="C")
            _ds.assert_called_once_with("Inflow", "Reference")
            _fc.assert_called_once_with(["A", "B"], foo="A", bar="C", squeeze=False)

@pytest.mark.parametrize("func", [stats.mannwhitneyu, stats.wilcoxon])
@pytest.mark.parametrize(
    ("x", "all_same"), [([5, 5, 5, 5, 5], True), ([5, 6, 7, 7, 8], False)]
)
def test_dist_compare_wrapper(x, all_same, func):
    y = [5, 5, 5, 5, 5]
    with mock.patch.object(stats, func.__name__) as _test:
        result = _dist_compare(x, y, _test)
        if all_same:
            assert numpy.isnan(result.stat)
            assert numpy.isnan(result.pvalue)
            assert _test.call_count == 0
        else:
            # assert result == (0, 0)
            _test.assert_called_once_with(x, y, alternative="two-sided")


#!/usr/bin/env python

import os
import glob
import numpy as np
import fitsio
import matplotlib.pyplot as plt
from matplotlib import ticker
from matplotlib.backends.backend_pdf import PdfPages
from astropy.table import Table

from bokpipe import *
from bokpipe.bokoscan import _convertfitsreg

def init_data_map(datadir,outdir,expTimes=None,files=None):
	dataMap = {}
	if not os.path.exists(outdir):
		os.mkdir(outdir)
	dataMap['outdir'] = outdir
	if files is None:
		dataMap['files'] = sorted(glob.glob(datadir+'*.fits') +
		                          glob.glob(datadir+'*.fits.gz') +
		                          glob.glob(datadir+'*.fits.fz'))
	else:
		dataMap['files'] = files
	dataMap['rawFiles'] = dataMap['files']
	dataMap['oscan'] = bokio.FileNameMap(outdir)
	dataMap['proc'] = bokio.FileNameMap(outdir,'_p')
	dataMap['files'] = [ dataMap['oscan'](f) for f in dataMap['files'] ]
	if expTimes is None:
		dataMap['expTime'] = np.array([fitsio.read_header(f)['EXPTIME']
		                                  for f in dataMap['files']])
	else:
		dataMap['expTime'] = expTimes
	try:
		# assume they are all the same
		dataMap['dataSec'] = \
		         _convertfitsreg(fitsio.read_header(
		                             dataMap['files'][0],'IM4')['DATASEC'])
	except IOError:
		pass
	return dataMap

def process_data(dataMap,redo=True,withvar=True,oscanims=False,bias2d=False):
	oscanSubtract = BokOverscanSubtract(output_map=dataMap['oscan'],
	                                    overwrite=redo,
		                                write_overscan_image=oscanims,
		                    oscan_cols_file=dataMap['outdir']+'oscan_cols',
		                    oscan_rows_file=dataMap['outdir']+'oscan_rows',
		                                verbose=10)#method='median_value')
	oscanSubtract.process_files(dataMap['rawFiles'])
	if bias2d:
		biasname = 'bias'
		biasStack = bokproc.BokBiasStack(#reject=None,
		                                 overwrite=redo,
		                                 with_variance=withvar)
		bias2dFile = os.path.join(dataMap['outdir'],biasname+'.fits')
		biasStack.stack(dataMap['biasFiles'],bias2dFile)
		#imProcess = bokproc.BokCCDProcess(bias2dFile,
		#                                  output_map=dataMap['proc'])
		#imProcess.process_files(flatFrames)

def imstat(dataMap,outfn='stats'):
	from astropy.stats import sigma_clip
	from scipy.stats import mode,scoreatpercentile
	array_stats = bokutil.array_stats
	fnlen = len(os.path.basename(dataMap['files'][0]))
	st = np.zeros(len(dataMap['flatSequence']),
	              dtype=[('file','S%d'%fnlen),
	                     ('expTime','f4'),
	                     ('median','16f4'),
	                     ('mean','16f4'),
	                     ('mode','16f4'),
	                     ('iqr25','16f4'),
	                     ('iqr75','16f4'),
	                     ('iqr10','16f4'),
	                     ('iqr90','16f4')])
	for _i,i in enumerate(dataMap['flatSequence']):
		expTime = dataMap['expTime'][i]
		fn = os.path.basename(dataMap['files'][i])
		fits = fitsio.FITS(dataMap['files'][i])
		print '%s %4.1f  ' % (fn,expTime),
		st['file'][_i] = fn
		st['expTime'][_i] = expTime
		for j,extn in enumerate(['IM%d' % n for n in range(1,17)]):
			modeVal,pix = array_stats(fits[extn].read()[dataMap['statsPix']],
			                          method='mode',retArray=True)
			st['mode'][_i,j] = modeVal
			st['mean'][_i,j] = pix.mean()
			st['median'][_i,j] = np.ma.median(pix)
			st['iqr25'][_i,j] = scoreatpercentile(pix,25)
			st['iqr75'][_i,j] = scoreatpercentile(pix,75)
			st['iqr10'][_i,j] = scoreatpercentile(pix,10)
			st['iqr90'][_i,j] = scoreatpercentile(pix,90)
			print '%5d ' % (modeVal),
		print
	fitsio.write(outfn+'.fits',st,clobber=True)

def scaled_histograms(dataMap,nims=None,outfn='pixhist'):
	pdf = PdfPages(outfn+'.pdf')
	for _i,i in enumerate(dataMap['flatSequence']):
		if nims is not None and _i==nims:
			break
		expTime = dataMap['expTime'][i]
		expScale = dataMap['refExpTime'] / expTime
		print dataMap['files'][i]
		fn = os.path.basename(dataMap['files'][i])
		fits = fitsio.FITS(dataMap['files'][i])
		fig = plt.figure(figsize=(8.0,10))
		plt.subplots_adjust(0.08,0.08,0.92,0.92,0.3,0.35)
		for j,extn in enumerate(['IM%d' % n for n in range(1,17)]):
			ax = plt.subplot(8,2,j+1)
			pix = fits[extn].read()[dataMap['statsPix']]
			ax.hist(expScale*pix.flatten(),100,(0,40000),edgecolor='none')
			ax.text(0.05,0.9,extn,va='top',size=9,transform=ax.transAxes)
			ax.set_xlim(0,40000)
			ax.xaxis.set_major_locator(ticker.MultipleLocator(10000))
			ax.xaxis.set_minor_locator(ticker.MultipleLocator(2000))
			ax.yaxis.set_major_locator(ticker.MultipleLocator(50000))
		plt.figtext(0.5,0.99,fn+' exp=%.1f' % expTime,ha='center',va='top')
		pdf.savefig(fig)
		plt.close(fig)
	pdf.close()

def plot_sequence(dataMap,st,imNum,which='median'):
	expScale = dataMap['refExpTime']/st['expTime']
	seqno = 1 + np.arange(len(st))
	ref = np.isclose(expScale,1.0)
	j = imNum - 1
	plt.figure(figsize=(8,6))
	plt.subplots_adjust(0.11,0.08,0.96,0.95)
	plt.errorbar(seqno[ref],expScale[ref]*st[which][ref,j],
	                   [expScale[ref]*(st[which]-st['iqr10'])[ref,j],
	                    expScale[ref]*(st['iqr90']-st[which])[ref,j]],
	             fmt='bs-')
	plt.errorbar(seqno[~ref],expScale[~ref]*st[which][~ref,j],
	                   [expScale[~ref]*(st[which]-st['iqr10'])[~ref,j],
	                    expScale[~ref]*(st['iqr90']-st[which])[~ref,j]],
	             fmt='cs-')
	#plt.scatter(seqno,expScale*st['mode'][:,j],marker='+',c='r')
	#plt.scatter(seqno,expScale*st['mean'][:,j],marker='x',c='g')
	plt.xlabel('sequence number')
	plt.ylabel('counts scaled by exp time')
	plt.title('IM%d'%imNum)
	plt.xlim(0.5,len(st)+0.5)

def fit_ref_exposures(dataMap,st,imNum,
                      which='median',method='spline',doplot=False):
	from scipy.interpolate import UnivariateSpline
	seqno = 1 + np.arange(len(st))
	t = st['expTime']
	ref = np.isclose(t,dataMap['refExpTime'])
	j = imNum - 1
	refCounts = st[which][ref,j][0]
	if method=='linear':
		_fit = np.polyfit(seqno[ref],refCounts/st[which][ref,j],1)
		fit = lambda x: np.polyval(_fit,x)
	elif method=='spline':
		fit = UnivariateSpline(seqno[ref],refCounts/st[which][ref,j],
		                       s=1e-5,k=3)
	else:
		raise ValueError
	if doplot:
		plt.figure()
		plt.subplot(211)
		plt.plot(seqno[ref],st[which][ref,j],'bs-')
		plt.plot(seqno,refCounts/fit(seqno),c='r')
		plt.subplot(212)
		plt.plot(seqno[ref],(st[which][ref,j]-refCounts/fit(seqno[ref]))
		                      /st[which][ref,j],'bs-')
		plt.axhline(0,c='r')
	return fit

def plot_linearity_curves(dataMap,st,which='median',correct=True,isPTC=False,
                          refCor=None,fitmethod='spline',outfn='linearity',
	                      onlyim=None):
	seqno = 1 + np.arange(len(st))
	t = st['expTime']
	print seqno,t
	refExpTime = dataMap['refExpTime']
	ref = np.isclose(t,refExpTime)
	refCorFit = None
	ii = np.arange(len(st))
	# only use the increasing sequence, not the reference exposures
	ii = ii[~ref]
	if isPTC:
		# for PTCs skip every other image since they are done in pairs
		ii = ii[::2]
	# only fit to unsaturated frames
	try:
		firstsat = np.where(np.any(st[which][ii,:] > 55000,axis=1))[0][0]
	except IndexError:
		firstsat = -1
	if onlyim is None:
		pdf = PdfPages(outfn+'.pdf')
	for imNum in range(1,17):
		if onlyim is not None and imNum != onlyim:
			continue
		j = imNum - 1
		# correct lamp variation
		if correct:
			if refCor is None:
				fscl_fit = fit_ref_exposures(dataMap,st,imNum,which,
				                             method=fitmethod)
			else:
				if refCorFit is None:
					refCorFit = fit_ref_exposures(dataMap,st,imNum,which)
				fscl_fit = refCorFit
			fscl = fscl_fit(seqno)
		else:
			fscl = np.ones_like(seqno)
		fit = np.polyfit(t[ii[:firstsat]],
		                 fscl[ii[:firstsat]]*st[which][ii[:firstsat],j],1)
		fitv = np.polyval(fit,t)
		slope = fit[0] / (st[which][ref,j][0]/refExpTime)
		#
		pltindex = imNum % 4
		if onlyim is None:
			if pltindex == 1:
				fig = plt.figure(figsize=(8,10))
				plt.subplots_adjust(0.11,0.08,0.96,0.95,0.25,0.2)
			ax = plt.subplot(4,2,2*(j%4)+1)
		else:
			fig = plt.figure(figsize=(6,2.5))
			plt.subplots_adjust(0.11,0.23,0.99,0.98,0.35,0.2)
			ax = plt.subplot(1,2,1)
		plt.plot(t[ii],fscl[ii]*st[which][ii,j],'bs-')
		plt.xlim(0.9*t.min(),t.max()+0.5)
		plt.xscale('log')
		plt.ylim(1e2,9e4)
		plt.yscale('log')
		plt.ylabel('counts [%s]' % which)
		tt = np.logspace(-1,np.log10(1.3*t.max()),100)
		plt.plot(tt,np.polyval(fit,tt),c='r')
		plt.text(0.05,0.9,'IM%d'%imNum,va='top',transform=ax.transAxes)
		plt.text(0.95,0.18,r'y = %.1f $\times$ t + %.1f' % tuple(fit),
		         ha='right',va='top',size=9,transform=ax.transAxes)
		plt.text(0.95,0.10,r'y = %.3f $\times$ counts + %.1f' % (slope,fit[1]),
		         ha='right',va='top',size=9,transform=ax.transAxes)
		if pltindex==0 or onlyim is not None:
			plt.xlabel('exptime (s)')
		#
		if onlyim is None:
			ax = plt.subplot(4,2,2*(j%4)+2)
		else:
			ax = plt.subplot(1,2,2)
		plt.plot(t[ii],100*(fscl[ii]*st[which][ii,j]-fitv[ii])/fitv[ii],'bs-')
		plt.axhline(0,c='r')
		#ax.xaxis.set_major_locator(ticker.MultipleLocator(10))
		#ax.xaxis.set_minor_locator(ticker.MultipleLocator(2))
		ax.yaxis.set_major_locator(ticker.MultipleLocator(2))
		ax.yaxis.set_minor_locator(ticker.MultipleLocator(0.5))
		plt.ylim(-5,5)
		plt.xlim(0.9*t.min(),t.max()+0.5)
		plt.xscale('log')
		if pltindex==0 or onlyim is not None:
			plt.xlabel('exptime (s)')
		plt.ylabel('residual \%')
		if onlyim is None:
			if pltindex == 0:
				pdf.savefig(fig)
				plt.close(fig)
	if onlyim is None:
		pdf.close()

def rel_gain(dataMap,st,which='median',correct=True,fitmethod='spline',
             nskip=0):
	seqno = 1 + np.arange(len(st))
	t = st['expTime']
	refExpTime = dataMap['refExpTime']
	ref = np.isclose(t,refExpTime)
	refCorFit = None
	ii = np.arange(len(st))
	ii = ii[~ref]
	ii = ii[nskip:]
	sky4 = st[which][ii,3]
	fit_ii = ii[np.where((sky4>5000)&(sky4<25000))[0]]
	plt.figure()
	for imNum in range(1,17):
		j = imNum - 1
		# correct lamp variation
		if correct:
			if True: #refCor is None:
				fscl_fit = fit_ref_exposures(dataMap,st,imNum,which,
				                             method=fitmethod)
			else:
				if refCorFit is None:
					refCorFit = fit_ref_exposures(dataMap,st,imNum,which)
				fscl_fit = refCorFit
			fscl = fscl_fit(seqno)
		else:
			fscl = np.ones_like(seqno)
		fit = np.polyfit(t[fit_ii],fscl[fit_ii]*st[which][fit_ii,j],1)
		fitv = np.polyval(fit,t)
#		slope = fit[0] / (st[which][ref,j][0]/refExpTime)
		xx = np.array(0,1.1*t.max())
		plt.subplot(4,4,imNum)
		if False:
			plt.scatter(t[ii],fscl[ii]*st[which][ii,j])
			plt.plot(xx,np.polyval(fit,xx),c='r')
		else:
			plt.scatter(t[ii],fscl[ii]*st[which][ii,j]/fitv[ii])
			plt.axhline(1,c='r')
		plt.ylim(0.7,1.3)
		if True:
			plt.xscale('log')
		plt.xlim(0.9*t.min(),1.1*t.max())

def get_first_saturated_frame(seq):
	try:
		firstsat = np.where(seq > 55000)[0][0]
	except IndexError:
		firstsat = -1
	return firstsat

def compare_oscan_levels(dataMap,st):
	files = [ dataMap['files'][i] for i in dataMap['flatSequence'] ]
	oscans = np.zeros((len(files),16))
	for j in range(16):
		oscans[:,j] = [ fitsio.read_header(f,'IM%d'%(j+1))['OSCANMED']
	                      for f in files ]
	seqno = 1 + np.arange(len(st))
	plt.figure()
	for j in range(8,16):
		ax = plt.subplot(8,2,2*(j%8)+1)
		i1 = get_first_saturated_frame(st['median'][:,j])
		plt.scatter(st['median'][:i1,j],oscans[:i1,j],c='b')
		plt.ylabel('IM%d'%(j+1))
		ax = plt.subplot(8,2,2*(j%8)+2)
		plt.scatter(seqno[:i1],oscans[:i1,j],c='b')

def init_sep09bss_data_map():
	datadir = os.environ.get('BASSDATA')+'/20150909/bss/20150908/'
	exptimes = np.loadtxt(datadir+'../bss.20150909.log',usecols=(3,))
	exptimes = exptimes[50:]
	print exptimes
	rdxdir = os.environ.get('GSCRATCH','tmp_sep')+'/bss_sep09/'
	if not os.path.exists(rdxdir):
		os.makedirs(rdxdir)
	dataMap = init_data_map(datadir,rdxdir,
	                        expTimes=exptimes,files=None)
	dataMap['rawFiles'] = dataMap['rawFiles'][50:]
	dataMap['files'] = dataMap['files'][50:]
	dataMap['biasFiles'] = dataMap['files'][-5:]
	#dataMap['flatSequence'] = range(50,68)
	dataMap['flatSequence'] = range(18)
	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
	dataMap['refExpTime'] = 40.0
	return dataMap

def init_sep29ptc_data_map():
	dataMap = init_data_map(
	      "/home/ian/dev/rapala/bokpipe/scratch/sep29ptcs/ptc/",'sep29ptcs/')
	dataMap['biasFiles'] = [dataMap['files'][0],]
	dataMap['flatSequence'] = range(1,len(dataMap['files']))
	dataMap['statsPix'] = np.s_[20:-20,100:-100]
	dataMap['refExpTime'] = 10.0
	return dataMap

def init_oct02ptc_data_map():
	dataMap = init_data_map(os.environ.get('GSCRATCH')+'/02oct15/ptc/',
	                        os.environ.get('GSCRATCH')+'/02oct15/ptc_proc/')
	dataMap['biasFiles'] = [dataMap['files'][0],]
	dataMap['flatSequence'] = range(1,len(dataMap['files']))
	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
	dataMap['refExpTime'] = 10.0
	return dataMap

def init_oct20_data_map():
	datadir = os.environ.get('BASSDATA')+'/20151020/'
	exptimes = np.loadtxt(datadir+'images.log',usecols=(6,))
	nuse = 53
	exptimes = exptimes[:nuse]
	print exptimes
	dataMap = init_data_map(datadir,'tmp_oct20',expTimes=exptimes)
	dataMap['rawFiles'] = dataMap['rawFiles'][:nuse]
	dataMap['files'] = [ dataMap['oscan'](f)
	                       for f in dataMap['files'][:nuse] ]
	dataMap['biasFiles'] = dataMap['files'][:20]
	dataMap['flatSequence'] = range(20,nuse)
	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
	dataMap['refExpTime'] = 3.0
	return dataMap

def init_nov11g_data_map():
	datadir = os.environ.get('BASSDATA')+'/Nov2015/'
	log = Table.read(datadir+'bassLog_Nov2015.fits')
	exptimes = log['expTime'][111:150]
	files = [ datadir+f['utDir']+'/'+f['fileName']+'.fits'
	              for f in log[111:150] ]
	dataMap = init_data_map(datadir,'tmp_nov11g',
	                        expTimes=exptimes,files=files)
	dataMap['biasFiles'] = dataMap['files'][-10:]
	dataMap['flatSequence'] = np.arange(len(dataMap['files'])-10)
	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
	dataMap['refExpTime'] = 3.0
	return dataMap

def init_nov14_data_map(filt):
	datadir = os.environ.get('BASSDATA')+'/Nov2015/'
	log = Table.read(datadir+'bassLog_Nov2015.fits')
	if filt=='g':
		frames = np.r_[np.s_[297:345],np.s_[247:257]]
	else:
		frames = np.r_[np.s_[345:393],np.s_[247:257]]
	exptimes = log['expTime'][frames]
	files = [ datadir+f['utDir']+'/'+f['fileName']+'.fits'
	              for f in log[frames] ]
	dataMap = init_data_map(datadir,'tmp_nov14'+filt,
	                        expTimes=exptimes,files=files)
	dataMap['biasFiles'] = dataMap['files'][-10:]
	dataMap['flatSequence'] = np.arange(len(dataMap['files'])-10)
	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
	dataMap['refExpTime'] = {'Ha':10.0,'g':3.0}[filt]
	return dataMap

def init_jan3_data_map(filt):
	datadir = os.environ.get('BASSDATA')
	log = Table.read('basslogs/log_ut20160103.fits')
	if filt=='g':
		frames = np.r_[np.s_[57:105],np.s_[160:170]]
	else:
		frames = np.r_[np.s_[105:160],np.s_[160:170]]
	exptimes = log['expTime'][frames]
	files = [ datadir+'/'+f['utDir'].strip()+'/'+f['fileName'].strip()+'.fits'
	              for f in log[frames] ]
	dataMap = init_data_map(datadir,'tmp_jan3'+filt,
	                        expTimes=exptimes,files=files)
	dataMap['biasFiles'] = dataMap['files'][-10:]
	dataMap['flatSequence'] = np.arange(len(dataMap['files'])-10)
	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
	dataMap['refExpTime'] = {'Ha':10.0,'g':3.0}[filt]
	return dataMap

def init_data_map_fromfile(filename,outdir='tmp',nersc=True):
	datadir = os.environ.get('BASSDATA')
	if nersc:
		datadir = os.path.join(datadir,'BOK_Raw')
	log = np.loadtxt(filename,dtype=[('frameNum','i4'),('utDir','S8'),
	                                 ('fileName','S35'),
	                                 ('imType','S10'),('filter','S8'),
	                                 ('expTime','f4')],skiprows=1)
	exptimes = log['expTime']
	files = [ datadir+'/'+f['utDir'].strip()+'/'+f['fileName'].strip()+'.fits'
	              for f in log ]
	if nersc:
		files = [ f+'.fz' for f in files ]
	dataMap = init_data_map(datadir,outdir,
	                        expTimes=exptimes,files=files)
	dataMap['biasFiles'] = np.array(dataMap['files'])[log['imType']=='zero']
	dataMap['flatSequence'] = np.where(log['imType']=='flat')[0]
	dataMap['statsPix'] = bokutil.stats_region('amp_corner_ccdcenter_small')
	# assume it starts with reference
	dataMap['refExpTime'] = exptimes[dataMap['flatSequence'][0]]
	return dataMap

if __name__=='__main__':
	import sys
	dataset = sys.argv[1]
	if dataset == 'sep09bss':
		dataMap = init_sep09bss_data_map()
	elif dataset == 'oct02':
		dataMap = init_oct02ptc_data_map()
	elif dataset == 'oct20':
		dataMap = init_oct20_data_map()
	elif dataset == 'nov11g':
		dataMap = init_nov11g_data_map()
	elif dataset == 'nov14g':
		dataMap = init_nov14_data_map('g')
	elif dataset == 'nov14Ha':
		dataMap = init_nov14_data_map('Ha')
	elif dataset == 'jan3g':
		dataMap = init_jan3_data_map('g')
	elif dataset == 'jan3Ha':
		dataMap = init_jan3_data_map('Ha')
	else:
		dataMap = init_data_map_fromfile(sys.argv[2],dataset)
	print 'processing ',dataset
	if not os.path.exists('stats_'+dataset+'.fits'):
		process_data(dataMap,bias2d=True)
		imstat(dataMap,outfn='stats_'+dataset)
	st = fitsio.read('stats_'+dataset+'.fits')
	plot_linearity_curves(dataMap,st,outfn='linearity_'+dataset)
	if True:
		plot_linearity_curves(dataMap,st,outfn='linearity_'+dataset,
		                      onlyim=4)
		plt.savefig('linearity_IM4_%s.png'%dataset)
		plot_sequence(dataMap,st,4)
		plt.savefig('linsequence_IM4_%s.png'%dataset)


#!/usr/bin/env python

import cv2
import sys
import yaml

import signal
import numpy as np
#import utm

import matplotlib as mpl
import matplotlib.cm as cm

import rospy

import argparse

import actionlib

from cosmos_msgs.msg import KrigInfo
from cosmos_msgs.srv import CompareModels
import kriging_exploration.map_coords
import std_msgs.msg

import open_nav.msg

from kriging_exploration.data_grid import DataGrid
from kriging_exploration.map_coords import MapCoords
from kriging_exploration.visualiser import KrigingVisualiser
from kriging_exploration.canvas import ViewerCanvas
from kriging_exploration.topological_map import TopoMap
from kriging_exploration.exploration import ExplorationPlan

from sensor_msgs.msg import NavSatFix

def overlay_image_alpha(img, img_overlay):
    """Overlay img_overlay on top of img at the position specified by
    pos and blend using alpha_mask.
    """
    show_image = img.copy()
    alpha =  img_overlay[:, :, 3] / 255.0   # Alpha mask must contain values
                                            # within the range [0, 1]
                                            # and be the same size as img_overlay.
    # Image ranges
    y1, y2 = 0, img.shape[0]
    x1, x2 = 0, img.shape[1]

    channels = img.shape[2]

    alpha_inv = 1.0 - alpha

    for c in range(channels):
        show_image[y1:y2, x1:x2, c] = (alpha * img_overlay[y1:y2, x1:x2, c] + alpha_inv * img[y1:y2, x1:x2, c])
    return show_image

class Explorator(KrigingVisualiser):

    #_w_shape=[(0, 16), (1, 17), (3, 17), (5, 16), (8, 15), (10, 15), (12, 14), (14, 13), (12, 12), (10, 11), (8, 11), (5, 10), (8, 9), (10, 9), (12, 8), (14, 7), (12, 6), (10, 5), (8, 5), (6, 4), (4, 3), (3, 2), (4, 1), (5, 0), (7, 0)]
    #_w_shape=[(17, 0), (17, 1), (17, 3), (16, 5), (15, 8), (15, 10), (14, 12), (13, 14), (12, 12), (11, 10), (11, 8), (10, 5), (9, 8), (9, 10), (8, 12), (7, 14), (6, 12), (5, 10), (5, 8), (4, 6), (3, 4), (2, 3), (1, 4), (0, 5), (0, 7)]
    #_w_shape=[(17, 0), (17,1), (17, 2), (17, 4), (16, 4), (16, 6), (16, 8), (15, 8), (15, 10), (14, 10), (14, 12), (13, 12), (13, 14), (12, 14), (12, 12), (11, 12), (11, 10), (10, 10), (10, 8), (10, 6), (10, 4), (9, 4), (9, 6), (9, 8), (9, 10), (8, 10), (8, 12), (7, 12), (7, 14), (6, 14), (6, 12), (5, 12), (5, 10), (4, 10), (4, 8), (4, 6), (4, 4), (3, 4), (3, 3), (2, 3), (2, 4), (1,4), (1, 6), (0,6), (1, 8), (0,8), (1, 10), (0, 10), (0, 12), (0, 14)]
    _w_shape=[(17, 0), (16, 1), (14, 6), (12, 11), (10, 14), (8, 9), (5, 14), (3, 11), (2, 6), (0, 3)]
    def __init__(self, lat_deg, lon_deg, zoom, size, args):
        self.targets = []
        self.results =[]
        self.result_counter=0
        self.explodist=0
        self.running = True
        self.last_coord=None
        signal.signal(signal.SIGINT, self.signal_handler)
        self.expid=args.experiment_name
        print "Creating visualiser object"
        super(Explorator, self).__init__(lat_deg, lon_deg, zoom, size)

        cv2.namedWindow('explorator')
        cv2.setMouseCallback('explorator', self.click_callback)

        self.current_model=-1
        self.draw_mode = 'none'
        self.grid = DataGrid(args.limits_file, args.cell_size)
        self.topo_map= TopoMap(self.grid)
        self.visited_wp=[]

        explo_type = args.area_coverage_type
        self.define_exploration_type(explo_type)

        self.navigating = False
        self.pause_exp = False
        self.exploring = 0
        self.n_inputs = 0

        print "NUMBER OF TARGETS:"
        print len(self.explo_plan.targets)

        self.limits_canvas = ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res)
        self.grid_canvas = ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res)
        self.exploration_canvas = ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res)
        self.gps_canvas = ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res)

        self.limits_canvas.draw_polygon(self.grid.limits, (0,0,255,128), thickness=1)
        self.grid_canvas.draw_grid(self.grid.cells, args.cell_size, (128,128,128,2), thickness=1)

        self.redraw()

        self.redraw_kriged=True
        self.redraw_var=True
        self.redraw_devi=True

        self.model_canvas=[]
        self.model_legend=[]
        self.kriging_canvas=[]
        self.klegend_canvas=[]
        self.klegend2_canvas=[]
        self.klegend3_canvas=[]
        self.sigma_canvas=[]
        self.sigma2_canvas=[]
        self.model_canvas_names=[]

        self.mean_out_canvas = ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res)
        self.mean_out_legend_canvas = ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res)
        self.mean_var_canvas = ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res)
        self.mean_var_legend_canvas = ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res)
        self.mean_dev_canvas = ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res)
        self.mean_dev_legend_canvas = ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res)

        rospy.loginfo("Subscribing to Krig Info")
        rospy.Subscriber("/kriging_data", KrigInfo, self.data_callback)
        rospy.Subscriber("/fix", NavSatFix, self.gps_callback)
        rospy.Subscriber('/penetrometer_scan', std_msgs.msg.String, self.scan_callback)
        self.req_data_pub = rospy.Publisher('/request_scan', std_msgs.msg.String, latch=False, queue_size=1)

        rospy.loginfo(" ... Connecting to Open_nav")

        self.open_nav_client = actionlib.SimpleActionClient('/open_nav', open_nav.msg.OpenNavAction)
        self.open_nav_client.wait_for_server()

        rospy.loginfo(" ... done")

        tim1 = rospy.Timer(rospy.Duration(0.2), self.drawing_timer_callback)
        tim2 = rospy.Timer(rospy.Duration(0.1), self.control_timer_callback)
        self.refresh()

        while(self.running):
            cv2.imshow('explorator', self.show_image)
            k = cv2.waitKey(20) & 0xFF
            self._change_mode(k)

        tim1.shutdown()
        tim2.shutdown()
        cv2.destroyAllWindows()
        sys.exit(0)

    # EXPLORATION PARAMS HERE!!!!
    def define_exploration_type(self, explo_type):
        self.exploration_strategy=explo_type
        self.n_goals=10

        if explo_type=='area_split':
            self.grid._split_area(3,3)
            sb=[]
            for i in self.grid.area_splits_coords:
                (y, x) = self.grid.get_cell_inds_from_coords(i)
                sb.append((x,y))
            self.explo_plan = ExplorationPlan(self.topo_map, args.initial_waypoint, args.initial_percent, ac_model=explo_type, ac_coords=sb)
        elif explo_type=='random':
            self.explo_plan = ExplorationPlan(self.topo_map, args.initial_waypoint, args.initial_percent)
        elif explo_type=='w_shape':
            self.explo_plan = ExplorationPlan(self.topo_map, args.initial_waypoint, args.initial_percent, ac_model=explo_type, ac_coords=self._w_shape)
        else: #greedy
            self.explo_plan = ExplorationPlan(self.topo_map, args.initial_waypoint, args.initial_percent, exploration_type='greedy', ac_model=explo_type)

    def drawing_timer_callback(self, event):
        self.refresh()

    def control_timer_callback(self, event):
        if self.navigating:
            if self.open_nav_client.simple_state ==2:
                print "DONE NAVIGATING"
                self.navigating = False
                if self.exploring==1:
                    self.exploring=2

        elif self.exploring==2:
            if not self.pause_exp:
                self.explo_plan.explored_wp.append(self.explo_plan.route.pop(0))
                info_str='Do_reading'
                self.req_data_pub.publish(info_str)
                self.exploring=3

        elif self.exploring==4:
            if not self.pause_exp:
                if len(self.explo_plan.route) >0:
                    gg=self.explo_plan.route[0]
                    self.open_nav_client.cancel_goal()
                    targ = open_nav.msg.OpenNavActionGoal()

                    targ.goal.coords.header.stamp=rospy.Time.now()
                    targ.goal.coords.latitude=gg.coord.lat
                    targ.goal.coords.longitude=gg.coord.lon

                    print "Going TO: ", gg
                    self.exploring=1
                    self.navigating=True
                    self.open_nav_client.send_goal(targ.goal)
                else:
                    print "Done Exploring"
                    self.exploring = 0
#        else:
#            if self.exploring:
#                print "waiting for new goal"

    def gps_callback(self, data):
        if not np.isnan(data.latitude):
            self.gps_canvas.clear_image()
            gps_coord = MapCoords(data.latitude,data.longitude)
            self.gps_canvas.draw_coordinate(gps_coord,'black',size=2, thickness=2, alpha=255)
            if self.last_coord:
                dist = gps_coord - self.last_coord
                self.explodist+= dist[0]
            self.last_coord=gps_coord

    def data_callback(self, msg):
        point_coord = kriging_exploration.map_coords.coord_from_satnav_fix(msg.coordinates)
        for i in msg.data:
            self.grid.add_data_point(i.model_name, point_coord, i.measurement)

        self.vmin, self.vmax = self.grid.get_max_min_vals()
        self.n_models=len(self.grid.models)

        for i in self.grid.models:
            if i.name not in self.model_canvas_names:
                print i.name
                self.model_canvas_names.append(i.name)
                self.model_canvas.append(ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res))
                self.model_legend.append(ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res))
                self.kriging_canvas.append(ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res))
                self.klegend_canvas.append(ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res))
                self.klegend2_canvas.append(ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res))
                self.klegend3_canvas.append(ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res))
                self.sigma_canvas.append(ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res))
                self.sigma2_canvas.append(ViewerCanvas(self.base_image.shape, self.satellite.centre, self.satellite.res))
            self.draw_inputs(self.model_canvas_names.index(i.name))

        self.n_inputs+=1
        if self.exploring==3:
            if self.n_inputs>3:
                self.krieg_all_mmodels()
                rospy.sleep(0.1)
                self.grid.calculate_mean_grid()
                rospy.sleep(0.1)
                self.draw_means()
                self.draw_mode="means"

                resp = self.get_errors()
                self.result_counter+=1
                d={}
                d['step']=self.result_counter
                d['id']=self.expid
                d['ns']=len(self.explo_plan.targets)
                d['coord']={}
                d['coord']['lat']=self.last_coord.lat
                d['coord']['lon']=self.last_coord.lon
                d['dist']=float(self.explodist)
                d['results']={}
                d['results']['groundtruth']=resp
                d['results']['var']={}
                d['results']['var']['mean']={}
                d['results']['var']['mean']['mean']= float(np.mean(self.grid.mean_variance))
                d['results']['var']['mean']['max']= float(np.max(self.grid.mean_variance))
                d['results']['var']['mean']['min']= float(np.min(self.grid.mean_variance))

#                d['results']['var']['std']['mean']= np.mean(self.grid.mean_deviation)
#                d['results']['var']['std']['max']= np.max(self.grid.mean_deviation)
#                d['results']['var']['std']['min']= np.min(self.grid.mean_deviation)

                means=[]
                maxs=[]
                mins=[]
                for i in range(self.n_models):
                    means.append(float(np.mean(self.grid.models[i].variance)))
                    maxs.append(float(np.max(self.grid.models[i].variance)))
                    mins.append(float(np.min(self.grid.models[i].variance)))

                d['results']['models']={}
                d['results']['models']['means']=means
                d['results']['models']['maxs']=maxs
                d['results']['models']['mins']=mins

                rospy.sleep(0.1)
                self.results.append(d)
                if self.exploration_strategy == 'greedy':
                    nwp = len(self.explo_plan.route) + len(self.explo_plan.explored_wp)
                    print nwp, " nodes in plan"
                    if nwp <= self.n_goals:
                        #THIS IS the ONE
                        #self.explo_plan.add_limited_greedy_goal(self.grid.mean_variance, self.last_coord)

                        self.explo_plan.add_greedy_goal(self.grid.mean_variance)

                        #self.explo_plan.add_montecarlo_goal(self.grid.mean_variance, self.last_coord)

                #self.draw_mode="deviation"
#                self.current_model=0
#                if self.redraw_devi:
#                    self.draw_all_devs()
                self.redraw()
                rospy.sleep(0.1)
            self.exploring=4

    def scan_callback(self, msg):
        if msg.data == 'Reading':
            print "GOT READING!!!"
            cx, cy = self.grid.get_cell_inds_from_coords(self.last_coord)
            if cx <0 or cy<0:
                print "Reading outside the grid"
            else:
                print 'Reading at: ', cx, cy
                for i in self.topo_map.waypoints:
                    if (cy,cx) == i.ind:
                        print 'Setting: ', i.name, i.coord, "as Visited"
                        i.visited= True
                        self.visited_wp.append(i)
                        self.grid_canvas.draw_waypoints(self.topo_map.waypoints, (0,255,0,2), thickness=1)
                        self.grid_canvas.draw_waypoints(self.visited_wp, (0,0,255,2), thickness=1)
                        self.redraw()

    def refresh(self):
        #self.show_image = self.image.copy()
        #self.show_image = cv2.addWeighted(self.gps_canvas.image, 0.7, self.image, 1.0, 0)
        #self.show_image = transparentOverlay(self.image, self.gps_canvas.image)
        self.show_image = overlay_image_alpha(self.image,self.gps_canvas.image)

    def redraw(self):
        self.image = cv2.addWeighted(self.grid_canvas.image, 0.5, self.base_image, 1.0, 0)
        self.image = cv2.addWeighted(self.limits_canvas.image, 0.75, self.image, 1.0, 0)
        self.image = cv2.addWeighted(self.exploration_canvas.image, 0.75, self.image, 1.0, 0)
        if self.draw_mode == "inputs" and self.current_model>=0 :
            self.image = cv2.addWeighted(self.model_canvas[self.current_model].image, 0.75, self.image, 1.0, 0)
            self.image = overlay_image_alpha(self.image, self.model_legend[self.current_model].image)

        if self.draw_mode == "kriging":# and self.current_model>=0 :
            self.image = cv2.addWeighted(self.kriging_canvas[self.current_model].image, 0.75, self.image, 1.0, 0)
            #self.image = cv2.addWeighted(self.klegend_canvas[self.current_model].image, 1.0, self.image, 1.0, 0)
            self.image = overlay_image_alpha(self.image, self.klegend_canvas[self.current_model].image)

        if self.draw_mode == "deviation":# and self.current_model>=0 :
            self.image = cv2.addWeighted(self.sigma_canvas[self.current_model].image, 0.75, self.image, 1.0, 0)
            #self.image = cv2.addWeighted(self.klegend3_canvas[self.current_model].image, 1.0, self.image, 1.0, 0)
            self.image = overlay_image_alpha(self.image, self.klegend3_canvas[self.current_model].image)

        if self.draw_mode == "variance":# and self.current_model>=0 :
            self.image = cv2.addWeighted(self.sigma2_canvas[self.current_model].image, 0.75, self.image, 1.0, 0)
            #self.image = cv2.addWeighted(self.klegend2_canvas[self.current_model].image, 1.0, self.image, 1.0, 0)
            self.image = overlay_image_alpha(self.image, self.klegend2_canvas[self.current_model].image)

        if self.draw_mode == "means":
            self.image = cv2.addWeighted(self.mean_dev_canvas.image, 0.75, self.image, 1.0, 0)
            #self.image = cv2.addWeighted(self.klegend2_canvas[self.current_model].image, 1.0, self.image, 1.0, 0)
            self.image = overlay_image_alpha(self.image, self.mean_dev_legend_canvas.image)

        self.show_image = self.image.copy()

    def click_callback(self, event, x, y, flags, param):

        if event == cv2.EVENT_RBUTTONDOWN:
            click_coord = self.satellite._pix2coord(x,y)
            cx, cy = self.grid.get_cell_inds_from_coords(click_coord)

            if cx <0 or cy<0:
                print "click outside the grid"
            else:
                print cx, cy

            for i in self.topo_map.waypoints:
                if (cy,cx) == i.ind:
                    print i.name, i.coord.easting, i.coord.northing
                    i.visited= True
                    self.visited_wp.append(i)
                    self.grid_canvas.draw_waypoints(self.topo_map.waypoints, (0,255,0,2), thickness=1)
                    self.grid_canvas.draw_waypoints(self.visited_wp, (0,0,255,2), thickness=1)
                    self.redraw()

        if event == cv2.EVENT_LBUTTONDOWN:
            click_coord = self.satellite._pix2coord(x,y)
            cx, cy = self.grid.get_cell_inds_from_coords(click_coord)

            if cx <0 or cy<0:
                print "click outside the grid"
            else:
                print cx, cy

            for i in self.topo_map.waypoints:
                if (cy,cx) == i.ind:
                    self.open_nav_client.cancel_goal()
                    targ = open_nav.msg.OpenNavActionGoal()

                    #goal.goal.goal.header.
                    targ.goal.coords.header.stamp=rospy.Time.now()
                    targ.goal.coords.latitude=i.coord.lat
                    targ.goal.coords.longitude=i.coord.lon

                    print targ
                    self.navigating=True
                    self.open_nav_client.send_goal(targ.goal)
                    #self.client.wait_for_result()
                    # Prints out the result of executing the action
                    #ps = self.client.get_result()
                    #print ps

    def draw_inputs(self, nm):

        minv =  self.grid.models[nm].lims[0]
        maxv =  self.grid.models[nm].lims[1]

        if (maxv-minv) <=1:
            maxv = maxv + 50
            minv = minv - 50

        norm = mpl.colors.Normalize(vmin=minv, vmax=maxv)
        cmap = cm.jet
        colmap = cm.ScalarMappable(norm=norm, cmap=cmap)

        self.model_canvas[nm].clear_image()
        self.model_legend[nm].clear_image()

        for i in self.grid.models[nm].orig_data:
            cell = self.grid.cells[i.y][i.x]
            a= colmap.to_rgba(int(i.value))
            b= (int(a[2]*255), int(a[1]*255), int(a[0]*255), int(a[3]*50))
            self.model_canvas[nm].draw_cell(cell, self.grid.cell_size, b, thickness=-1)
            self.model_canvas[nm].put_text(self.grid.models[nm].name)

        self.model_legend[nm].put_text(self.grid.models[nm].name)
        self.model_legend[nm].draw_legend(minv, maxv, colmap, title="Kriging")

    def draw_krigged(self, nm):
        print "drawing kriging" + str(nm)

        minv =  self.grid.models[nm].min_val
        maxv =  self.grid.models[nm].max_val

        if (maxv-minv) <=1:
            maxv = maxv + 50
            minv = minv - 50

        norm = mpl.colors.Normalize(vmin=minv, vmax=maxv)
        cmap = cm.jet
        colmap = cm.ScalarMappable(norm=norm, cmap=cmap)

        self.kriging_canvas[nm].clear_image()
        self.klegend_canvas[nm].clear_image()
        for i in range(self.grid.models[nm].shape[0]):
            for j in range(self.grid.models[nm].shape[1]):
                cell = self.grid.cells[i][j]
                a= colmap.to_rgba(int(self.grid.models[nm].output[i][j]))
                b= (int(a[2]*255), int(a[1]*255), int(a[0]*255), int(a[3]*50))
                self.kriging_canvas[nm].draw_cell(cell, self.grid.cell_size, b, thickness=-1)

        self.klegend_canvas[nm].put_text(self.grid.models[nm].name)
        self.klegend_canvas[nm].draw_legend(minv, maxv, colmap, title="Kriging")

        self.redraw()

    def draw_variance(self, nm):
        print "drawing variance" + str(nm)

        minv =  self.grid.models[nm].min_var
        maxv =  self.grid.models[nm].max_var

        if (maxv-minv) <=1:
            maxv = maxv + 50
            minv = minv - 50

        norm = mpl.colors.Normalize(vmin=minv, vmax= maxv)
        cmap = cm.jet
        colmap = cm.ScalarMappable(norm=norm, cmap=cmap)

        self.sigma_canvas[nm].clear_image()
        self.klegend2_canvas[nm].clear_image()

        for i in range(self.grid.models[nm].shape[0]):
            for j in range(self.grid.models[nm].shape[1]):
                cell = self.grid.cells[i][j]
                a= colmap.to_rgba(int(self.grid.models[nm].variance[i][j]))
                b= (int(a[2]*255), int(a[1]*255), int(a[0]*255), int(a[3]*50))
                self.sigma2_canvas[nm].draw_cell(cell, self.grid.cell_size, b, thickness=-1)

        self.klegend2_canvas[nm].put_text(self.grid.models[nm].name)
        self.klegend2_canvas[nm].draw_legend(minv, maxv, colmap, title="Variance")
        self.redraw()

    def draw_means(self):
        print "drawing mean deviation ..."

        minv =  self.grid.min_mean_deviation
        maxv =  self.grid.max_mean_deviation

        if (maxv-minv) <=1:
            maxv = maxv + 50
            minv = minv - 50

        norm = mpl.colors.Normalize(vmin=minv, vmax=maxv)
        cmap = cm.jet
        colmap = cm.ScalarMappable(norm=norm, cmap=cmap)

        self.mean_dev_canvas.clear_image()
        self.mean_dev_legend_canvas.clear_image()

        for i in range(self.grid.shape[0]):
            for j in range(self.grid.shape[1]):
                cell = self.grid.cells[i][j]
                a= colmap.to_rgba(int(self.grid.mean_deviation[i][j]))
                b= (int(a[2]*255), int(a[1]*255), int(a[0]*255), int(a[3]*50))
                self.mean_dev_canvas.draw_cell(cell, self.grid.cell_size, b, thickness=-1)

        #self.mean_dev_legend_canvas.put_text(self.grid.models[nm].name)
        self.mean_dev_legend_canvas.draw_legend(minv, maxv, colmap, title="Mean Deviation")

        #self.draw_mode="means"
        self.redraw()

    def draw_deviation(self, nm):
        print "drawing deviation" + str(nm)

        minv =  self.grid.models[nm].min_dev
        maxv =  self.grid.models[nm].max_dev

        if (maxv-minv) <=1:
            maxv = maxv + 50
            minv = minv - 50

        norm = mpl.colors.Normalize(vmin=minv, vmax=maxv)
        cmap = cm.jet
        colmap = cm.ScalarMappable(norm=norm, cmap=cmap)

        self.sigma_canvas[nm].clear_image()
        self.klegend3_canvas[nm].clear_image()

        for i in range(self.grid.models[nm].shape[0]):
            for j in range(self.grid.models[nm].shape[1]):
                cell = self.grid.cells[i][j]
                a= colmap.to_rgba(int(self.grid.models[nm].deviation[i][j]))
                b= (int(a[2]*255), int(a[1]*255), int(a[0]*255), int(a[3]*50))
                self.sigma_canvas[nm].draw_cell(cell, self.grid.cell_size, b, thickness=-1)

        self.klegend3_canvas[nm].put_text(self.grid.models[nm].name)
        self.klegend3_canvas[nm].draw_legend(minv, maxv, colmap, title="Deviation")
        self.redraw()

    def krieg_all_mmodels(self):
        for i in self.grid.models:
            i.do_krigging()
            self.redraw_kriged=True
            self.redraw_var=True
            self.redraw_devi=True

    def draw_all_outputs(self):
        for i in self.grid.models:
            self.draw_krigged(self.model_canvas_names.index(i.name))
            self.redraw_kriged=False

    def draw_all_vars(self):
        for i in self.grid.models:
            self.draw_variance(self.model_canvas_names.index(i.name))
            self.redraw_var=False

    def draw_all_devs(self):
        for i in self.grid.models:
            self.draw_deviation(self.model_canvas_names.index(i.name))
            self.redraw_devi=False

    def _change_mode(self, k):
        if k == 27:
            self.running = False
        elif k == ord('q'):
            self.running = False
        elif k == ord('n'):
            print len(self.grid.models)
        elif k == ord('i'):
            if self.n_models > 0:
                self.draw_mode="inputs"
                self.current_model=0
                self.redraw()
        elif k == ord('d'):
            if self.n_models > 0:
                self.draw_mode="deviation"
                self.current_model=0
                if self.redraw_devi:
                    self.draw_all_devs()
                self.redraw()
        elif k == ord('v'):
            if self.n_models > 0:
                self.draw_mode="variance"
                self.current_model=0
                if self.redraw_var:
                    self.draw_all_vars()
                self.redraw()
        elif k == ord('t'):
            self.krieg_all_mmodels()
            self.grid.calculate_mean_grid()
            if self.n_models > 0:
                self.draw_all_outputs()
                self.draw_mode="kriging"
                self.current_model=0
                self.redraw()

        elif k == ord('k'):
            if self.n_models > 0:
                self.draw_mode="kriging"
                self.current_model=0
                if self.redraw_kriged:
                    self.draw_all_outputs()
                self.redraw()

        elif k == ord('>'):
            self.current_model+=1
            if self.current_model >= self.n_models:
                self.current_model=0
            self.redraw()
        elif k == ord('<'):
            self.current_model-=1
            if self.current_model < 0:
                self.current_model=self.n_models-1
            self.redraw()
        elif k == ord('w'):
            self.grid_canvas.draw_waypoints(self.topo_map.waypoints, (0,255,0,2), thickness=1)
            self.grid_canvas.draw_waypoints(self.visited_wp, (0,0,255,2), thickness=1)
            self.redraw()
        elif k == ord('e'):
            self.exploration_canvas.draw_waypoints(self.explo_plan.targets, (255,200,128,255), thickness=3)
            self.exploration_canvas.draw_plan(self.explo_plan.route, 'cyan', thickness=1)
            self.redraw()
            #xnames = [x.name for x in self.explo_plan.route]
            #print xnames
        elif k == ord('g'):
            if len(self.explo_plan.route) >0:
                gg=self.explo_plan.route[0]
                self.open_nav_client.cancel_goal()
                targ = open_nav.msg.OpenNavActionGoal()

                targ.goal.coords.header.stamp=rospy.Time.now()
                targ.goal.coords.latitude=gg.coord.lat
                targ.goal.coords.longitude=gg.coord.lon

                print "Going TO: ", gg
                self.exploring=1
                self.navigating=True
                self.open_nav_client.send_goal(targ.goal)
                self.result_counter=0
                self.explodist=0
            else:
                print "Done Exploring"
                self.exploring = 0
        elif k == ord('y'):
            vwp = []
            for i in self.visited_wp:
                vwp.append(i.name)
            yml = yaml.safe_dump(vwp, default_flow_style=False)
            fh = open("visited.yaml", "w")
            s_output = str(yml)
            fh.write(s_output)
            fh.close
        elif k == ord('l'):
            print "loading visited"

            with open("visited.yaml", 'r') as f:
                visited = yaml.load(f)
                for i in visited:
                    for l in self.topo_map.waypoints:
                        if i == l.name:
                            self.visited_wp.append(l)
                            break

        elif k == ord('a'):
            self.grid.calculate_mean_grid()
            self.draw_means()
            self.draw_mode="means"

        elif k == ord('p'):
            self.pause_exp= not self.pause_exp

        elif k == ord('c'):
            print self.grid.limits
            print "Area: ", self.grid.calculate_area(self.grid.limits)
            print "Area of Area: ", self.grid.area.area_size
            colours=['magenta','cyan', 'grey','white','red','yellow','green','blue']

            nc=0
            for j in self.grid.area_splits:
                print j.area_size
                #self.limits_canvas.draw_coordinate(j.centre, 'crimson', size=3, thickness=2)
                for i in j.limit_lines:
                    #self.limits_canvas.draw_line(i, colours[nc], thickness=1)
                    self.limits_canvas.draw_line(i, 'white', thickness=1)
                if nc < len(colours)-1:
                    nc+=1
                else:
                    nc=0

            self.redraw()

        elif k== ord('r'):
            #diff = (self.grid.models[1].output - self.grid.models[0].output)
            #print np.mean(diff), np.std(diff), diff.dtype
            print self.get_errors()

        elif k== ord('o'):
            print self.results
            outfile = self.expid + '.yaml'
            #print self.data_out
            yml = yaml.safe_dump(self.results, default_flow_style=False)
            fh = open(outfile, "w")
            s_output = str(yml)
            #print s_output
            fh.write(s_output)
            fh.close

    def get_errors(self):
        error_chain=[]
        shapeo = self.grid.models[0].output.shape

        #print vals
        print "Waiting for Service"
        rospy.wait_for_service('/compare_model')
        compare_serv = rospy.ServiceProxy('/compare_model', CompareModels)

        for i in range(self.n_models):
            try:
                d={}
                print "going for it ", i
                vals = np.reshape(self.grid.models[i].output, -1)
                resp1 = compare_serv('kriging', i, shapeo[0], shapeo[1], vals.tolist())
                d['name']= self.grid.models[i].name
                d['type']= 'kriging'
                d['errors']={}
                d['errors']['error']=resp1.error
                d['errors']['mse']=resp1.mse
                d['errors']['std']=resp1.std
                d['errors']['var']=resp1.var
                #print resp1
                error_chain.append(d)
            except rospy.ServiceException, e:
                print "Service call failed: %s"%e

        try:
            d={}
            print "Mean "
            vals = np.reshape(self.grid.mean_output, -1)
            resp1 = compare_serv('mean', 0, shapeo[0], shapeo[1], vals.tolist())
            #print self.grid.mean_output
            d['name']= 'mean'
            d['type']= 'mean'
            d['errors']={}
            d['errors']['error']=resp1.error
            d['errors']['mse']=resp1.mse
            d['errors']['std']=resp1.std
            d['errors']['var']=resp1.var

            #print resp1
            error_chain.append(d)
        except rospy.ServiceException, e:
            print "Service call failed: %s"%e

        return error_chain

    def signal_handler(self, signal, frame):
        self.running = False
        print('You pressed Ctrl+C!')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--cell_size", type=int, default=10,
                        help="cell size in meters")
    parser.add_argument("--initial_percent", type=float, default=0.05,
                        help="Percentage of cells to be explored on the initial plan")
    parser.add_argument("--limits_file", type=str, default='limits.coords',
                        help="Percentage of cells to be explored on the initial plan")
    parser.add_argument("--initial_waypoint", type=str, default='WayPoint498',
                        help="Percentage of cells to be explored on the initial plan")
    parser.add_argument("--area_coverage_type", type=str, default='area_split',
                        help="Type of area coverage, random or area_split")
    parser.add_argument("--experiment_name", type=str, default='exp1',
                        help="Experiment ID")
    args = parser.parse_args()

    rospy.init_node('kriging_exploration')
    #Explorator(53.261685, -0.527158, 16, 640, args.cell_size)

    #Explorator(53.267213, -0.533420, 17, 640, args)  #Football Field
    Explorator(53.261576, -0.526648, 17, 640, args)  #Half cosmos field
    #Explorator(53.261685, -0.525158, 17, 640, args) #COSMOS Field