# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.4.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Compare results in the DB

# %%
## To autoreload codein python files here.
# %load_ext autoreload
# %autoreload 2

## Auto-format cells to ease diffs.
# %load_ext lab_black

# %%
# %matplotlib ipympl

# %%
from typing import Union, List, Callable, Any, Sequence as Seq
import io, logging, re, sys
from pathlib import Path, PurePosixPath as P

from columnize import columnize
import numpy as np
import pandas as pd
from pandas import HDFStore
from pandas.core.generic import NDFrame
from matplotlib import pyplot as plt
import qgrid
import wltp
from wltp import io as wio, cycler
from wltp.experiment import Experiment

## Add tests/ into `sys.path` to import `vehdb` module.
#
proj_dir = str(Path(wltp.__file__).parents[1] / "tests")
if proj_dir not in sys.path:
    sys.path.insert(0, proj_dir)

import vehdb

idx = pd.IndexSlice
log = logging.getLogger("CarsDB-compare.ipynb")
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s|%(levelname)4.4s|%(module)s:[%(funcName)s]:\n  +--> %(message)s",
    datefmt="%Y-%m-%d,%H:%M:%S",
)

pd.set_option("display.max_columns", 32)

# %%
## DEFINITIONS
#
inp_h5fname = "VehData/WltpGS-msaccess.h5"
out_h5fname = "VehData/WltpGS-pyalgo.h5"
c_n, c_p, c_n_norm, c_p_norm = "n", "Pwot", "n_norm", "p_norm"

# %%
vehdb.print_nodes(inp_h5fname)
vehdb.print_nodes(out_h5fname)

# %%
from wltp.invariants import vround, nround1

def load_accdb_and_python_datasets(veh_nums=None):
    p1, c1 = vehdb.merge_db_vehicle_subgroups(
        inp_h5fname, "prop", "cycle", veh_nums=veh_nums
    )
    p2, c2 = vehdb.merge_db_vehicle_subgroups(
        out_h5fname, "oprop", "cycle", veh_nums=veh_nums
    )

    ## Originally fetched as 2-levels (veh, item) MultiIndex Series.
    p1 = p1.unstack()
    p2 = p2.unstack()

    ## By the spec, V rounded to 2-digits,
    #  But exporting MSAccess --> Excel outputs garbage decimals!
    #
    v_cols = "v v_orig v_cap v_downscale".split()
    c1[v_cols] = vround(c1[v_cols])

    ## accdb does not offer `n_max`.(?)
    p1["n_max"] = nround1(p1["n_max1 n_max2 n_max3".split()].max(axis=1))

    return p1, c1, p2, c2

p1, c1, p2, c2 = load_accdb_and_python_datasets()

# %%
## EXPORT RESULTS to upload them for the GS-group when releasing.
#
# p2.drop('v116', axis=0).to_excel('pyalgo-props-124cases-1.0.0.dev12.xlsx')
# c2.drop('v116', axis=0).to_excel('pyalgo-cycles-124cases-1.0.0.dev12.xlsx')

# %%
# Available PROPs
print(columnize(list(p1.columns), displaywidth=160))
print(columnize(list(p2.columns), displaywidth=160))

# %%
display(
    vehdb.grid(p1, fitcols=False),
    vehdb.grid(p2, fitcols=False),
    # vehdb.grid(c1, fitcols=0),
)

# %%
sr_cmpr = vehdb.Comparator(lambda d, c: d[:, c], no_styling=True)
dataset_names = "accdb Python".split()  # Must sort with "diff" column.

# %%
## Report PROP differences
#
#     ACCDB,  PYALGO
equivalent_columns = [
    ("Description", None),
    ("test_mass", None),
    ("kerb_mass", None),
    ("vehicle_class", "wltc_class"),
    # ("pmr_km", "pmr"),
    ("f_dsc_req", "f_dsc"),
    ("v_max", "v_max"),
    ("n_max1", "n95_high"),
    # ("???", "is_n_lim_vmax")
]

cdf = sr_cmpr.compare((p1.stack(), p2.stack()), equivalent_columns, dataset_names)
## Workaround qgrid's hate for hierarchical-columns:
#  https://github.com/quantopian/qgrid/issues/18#issuecomment-149321165
cdf.columns = [" ".join(col).strip() for col in cdf.columns.values]

display(vehdb.grid(cdf, fitcols=False, cwidth=100))
# with pd.option_context('max_rows', 130):
#     display(cdf)

# %%
# Available CYCLE columns
print(columnize(list(c1.columns), displaywidth=160))
print(columnize(list(c2.columns), displaywidth=160))

# %%
cmpr = vehdb.Comparator(lambda d, c: d.loc[idx[:, c]])

# %%
## Report CYCLE-MEAN differences
#
# # Vehicles with DOWNSCALE discrepancies
# veh_nums = [7, 19, 20, 33, 35, 43, 44, 56, 59, 66]
# veh_nums += [82, 88, 91, 99, 100, 101, 112, 113, 114]
# # UNCOMMENT next line to FETCH new vehicles.
p1, c1, p2, c2 = load_accdb_and_python_datasets()
equivalent_series = [
    ("v_orig", "V_cycle"),
    ("v_downscale", "v_target"),
    # ("a", "a_target"),
    ("P_tot_set", "p_req"),
    # ("P_max", "p_available"),
    ("g_max", "g_max0"),
    ("g_min", "g_min"),
    # ("gear", "gears"),
    # ("nc", "n"),
]

c2.columns = wio.flatten_columns(c2.columns)

cols1, cols2 = zip(*equivalent_series)
cols1, cols2 = list(cols1), list(cols2)
## Concat props to col-aggregates & convert prop-cols to numerics
cc1 = pd.concat((p1.infer_objects(), c1[cols1].abs().mean(level=0)), axis=1)
cc2 = pd.concat((p2.infer_objects(), c2[cols2].abs().mean(level=0)), axis=1)

equivalent_props = [
    ("Description", None),
    ("vehicle_class", "wltc_class"),
    ("pmr_km", None),
    ("no_of_gears", None),
    ("f_dsc_req", "f_dsc"),
    ("v_max", "v_max"),
    ("n_vmax", "n_vmax"),
    ("gear_v_max", "g_vmax"),
    ("n_max1", "n_max1"),
    ("n_max2", "n_max2"),
    ("n_max3", "n_max3"),
    ("n_max", "n_max"),
]
display(
    cmpr.compare(
        (cc1, cc2), equivalent_props + equivalent_series, dataset_names, describe=True
    )
)

# %%
## Repeat, to compare while coding.
# display(cmpr.compare((cc1, cc2), equivalent_props + equivalent_series, dataset_names, describe=True))

# %% [markdown]
# AccDB vehicles: 42, 46, 52, 53 & 90 have broken `v_max`, 48 has broken `wot(ASM)`.

# %%
known_bads = set(wio.veh_names([42, 46, 48, 52, 53, 90]))
display(
    cmpr.compare(
        (cc1, cc2), equivalent_props + equivalent_series, dataset_names
    ).set_properties(subset=(known_bads, idx[:]), color="red")
)

# %%
## Repeat, to compare while coding.
# display(cmpr.compare((cc1, cc2), equivalent_props + equivalent_series, dataset_names).set_properties(subset=(known_bads, idx[:]), color='red'))

# %% [markdown]
# ## Compare a vehicle from AccDB <-->PyAlgo *interactively*
# **TODO:** collect and hide all this comparison GUI code below into a python module.

# %%
case_loaded = [None, None]

def _load_interactive_case(
    case_name,
    # increase vertical seperation between flags (but do not exceed max(V))
    flag_mul=2,
):
    accdb_cycle = c1.loc[case_name].copy()
    accdb_gears = [c for c in accdb_cycle.columns if c.startswith("Ind_g")]

    cycle = c2.loc[case_name].dropna(how="all", axis=1).copy()

    ## Scale each flag into a different value, to plot separately, and
    #  to plot in the same axis as V
    #  (bc when plotting flags in `secondary_y`, grid is not working)
    #

    all_flags = [c for c in cycle.columns if c.startswith("ok_")]
    ok_flags = [c for c in all_flags if not c.startswith("ok_gear/")]
    flag_count = len(ok_flags)
    ok_flags = cycle.loc[:, ok_flags].copy()

    ok_gears = [c for c in all_flags if c.startswith("ok_gear/")]
    gear_count = len(ok_gears)
    ok_gears = cycle.loc[:, ok_gears]

    ok_flags[ok_flags < 0] = np.NAN  # Restore NANFLAG --> NAN
    ok_flags = ok_flags * (np.arange(flag_count) + 1) * flag_mul
    ok_gears = ok_gears * (np.arange(gear_count) + flag_count + 1) * flag_mul

    cycle.columns = wio.inflate_columns(cycle.columns)
    ok_flags.columns = wio.inflate_columns(ok_flags.columns)
    ok_gears.columns = wio.inflate_columns(ok_gears.columns)

    return cycle, ok_flags, ok_gears, accdb_cycle, accdb_gears, p1.loc[case_name]

def load_interactive_case(
    case_name,
    # increase vertical seperation between flags (but do not exceed max(V))
    flag_mul=2,
):
    if case_loaded[0] == case_name:
        cycle_data = case_loaded[1]
    else:
        cycle_data = _load_interactive_case(case_name)
        case_loaded[0] = case_name
        case_loaded[1] = cycle_data
    return cycle_data

# cycle, ok_flags, ok_gears, accdb_cycle, accdb_gears, accdb_props = load_interactive_case('v001')
# ok_flags.columns, ok_gears.columns

# %%
def decide_signal_axis(colnames):
    def is_velocity(col):
        return is_velocity(col[0]) if isinstance(col, tuple) else col.startswith("v_")

    l1 = [i for i in colnames if is_velocity(i)]
    l2 = [i for i in colnames if not is_velocity(i)]
    return l1, l2

def define_schemes():
    """Predfined case/pan/zooms of the compraison diagram below."""
    return [
        ## (label, caseno, zoom, pan, *other-UNUSED)
        ("dive(t=764): 1->2 gear too early", "v001", 45.0, 30.40),
        ("dive(t=903): 1st gear not reached", "v001", 45.0, 36.4),
        ("lowPower(t=724)", "v001", 43.0, 69.6),
        ("Py:LowP AccDB:DnShift(t=1571-8)", "v019", 69, 63.4),
        ("p_max from lower g(t=1571-8)", "v020", 69, 63.4),
        ("sameCar?", "v023", 1, 0),
        ("noPowerDives(t=1540-80)", "v024", 26, 63.60),
        ("noPowerDive(t=1574)", "v024", 26, 66.4),
        ("noPowerDive(t=1574)", "v025", 69, 63.4),
        (
            "decelToStop(t=1591)",
            "v025",
            69,
            72.0,
            "AccDB not respecting n_min=0.9 x n_idle (Annex 2-3.k.3)",
        ),
        ("lowPowerDive", "v033", 5, 72),
        ("low -1 gears?", "v035", 65, 63, "Diffs in gears above 2"),
        ("noPowerDive(t=1574)", "v024", 26, 66.4),
        ("g5Ok?(1542)", "v075", 55, 62.4, "Vehicle has too many insufficient powers"),
        ("lowP, lowG betterPAvail(t=1574)", "v087", 65, 63.40),
        ("lowPower(t=1574)", "v088", 65, 63.40),
        ("lowPower", "v101", 1, 0),
        ("lowPower", "v111", 1, 0),
        ("extention", "v117", 13, 72),
        ("extention", "v118", 13, 72),
        ("extention", "v119", 13, 72),
        ("extention", "v120", 13, 72),
        ("extention", "v121", 13, 72),
        ("lowPower", "v124", 1, 0),
        ("lowPower", "v125", 1, 0),
    ]

def grid(df):
    """Display a dynamic grid if given `df` too long. """
    display(vehdb.grid(df, fitcols=len(df.shape) < 2 or df.shape[1] < 12))

def merge_pyalgo_accdb(pyalgo, accdb):
    pyalgo = pyalgo.copy()
    pyalgo.columns = wio.flatten_columns(pyalgo.columns)
    merged = pd.concat((pyalgo, accdb), axis=1, keys=("pyalgo", "accdb"))
    merged.columns = wio.flatten_columns(merged.columns)

    return merged

def display_diff_gears(pyalgo, accdb, props):
    merged = merge_pyalgo_accdb(pyalgo, accdb)
    diffgears_idx = merged["pyalgo/g_max0"] != merged["accdb/g_max"]
    display(merged.loc[diffgears_idx, ["pyalgo/g_max0", "accdb/g_max"]])

def display_diff_cycle(pyalgo, accdb, props):
    merged = merge_pyalgo_accdb(pyalgo, accdb)
    diffgears_idx = merged["pyalgo/g_max0"] != merged["accdb/g_max"]
    grid(merged[diffgears_idx])

def display_pyalgo(pyalgo, accdb, props):
    pyalgo.columns = wio.flatten_columns(pyalgo.columns)
    grid(pyalgo.drop("t", axis=1))

def display_pyalgo_flags(pyalgo, accdb, props):
    flags = pyalgo.select_dtypes("int8").copy()
    flags.columns = wio.flatten_columns(flags.columns)
    pyalgo = pyalgo[["p_req", "p_avail"]].copy()
    pyalgo.columns = wio.flatten_columns(pyalgo.columns)
    grid(pd.concat((pyalgo, flags), axis=1))

def display_accdb_cycle(pyalgo, accdb, props):
    grid(accdb)

def display_accdb_props(pyalgo, accdb, props):
    grid(props)

## What to display in the Tabs beneath the plot
#
out_specs = [
    (display_diff_gears, "DiffGears"),
    (display_diff_cycle, "DiffGear cycles"),
    (display_pyalgo_flags, "PyAlgo flags"),
    (display_pyalgo, "PyAlgo"),
    (display_accdb_cycle, "AccDb cycle"),
    (display_accdb_props, "AccDb props"),
]

# %%
## TODO: CLASSify code in these cells.
from ipywidgets import (
    interact,
    interactive,
    interactive_output,
    fixed,
    interact_manual,
    widgets,
)
from IPython.display import clear_output

init_zoom = 55.0
init_pan = 30.40
max_zoom = 72.0

Case = widgets.SelectionSlider(options=list(c2.index.levels[0]), description="Case")
Gear = widgets.SelectionSlider(options=["g1"], description="Gear")
Zoom = widgets.FloatSlider(
    init_zoom, min=1.0, max=max_zoom, step=4.0, description="Zoom"
)
Pan = widgets.FloatSlider(init_pan, min=0.0, max=max_zoom, step=0.4, description="Pan")
IsPyAlgoGears = widgets.Checkbox(False, description="Plot PyAlgo Gear flags?")
IsAccdbGears = widgets.Checkbox(False, description="Plot AccDB Gear flags?")
PyAlgoSignals = widgets.SelectMultiple(rows=7, description="Pyalgo signals")
AccDBSignals = widgets.SelectMultiple(rows=7, description="AccDB signals")
AxisScenes = widgets.Select(options=[], description="Scenes")
AxisScenes.layout.width = "auto"
Desc = widgets.Textarea(disabled=True)  # DEFUNCT
Desc.layout.width = "24"
Desc.layout.height = "7em"
Selections = widgets.HBox(
    [
        Case,
        Zoom,
        Pan,
        AxisScenes,
        IsPyAlgoGears,
        PyAlgoSignals,
        IsAccdbGears,
        AccDBSignals,
        Gear,
    ]
)
Selections.layout.width = "95%"
Selections.layout.flex_flow = "row wrap"
# Selections.layout.justify_content = "flex-start"

#: Updated by *interact* function, for display functions to read them
#: pyalgo, accdb, accdb_props
results = [None, None, None]

Tab = widgets.Tab()

def refresh_tabs():
    global out_specs

    AxisScenes.options = [
        (f"{caseno}: {label}", (caseno, zoom, pan, *other))
        for label, caseno, zoom, pan, *other in define_schemes()
    ]
    needs_retabbing = False
    if out_specs and len(out_specs[0]) < 3:
        out_specs = [(widgets.Output(), *i) for i in out_specs]
        needs_retabbing = True
    if needs_retabbing or not Tab.children:
        Tab.children = [out for out, _, _ in out_specs]
        for i, (_, _, title) in enumerate(out_specs):
            Tab.set_title(i, title)

    return out_specs

refresh_tabs()

def update_tab_contents(change):
    if change.new is None:
        return
    out, func, _ = out_specs[change.new]
    has_content = bool(out.get_state()["outputs"])
    if not has_content:
        with out:
            func(*(i.copy() for i in results))

Tab.observe(update_tab_contents, names="selected_index")

Gui = widgets.VBox([Selections, Tab])

def update_valid_signals(change):
    case_name = Case.value
    pyalgo, _, ok_gears, accdb, _, _ = load_interactive_case(case_name)

    Gear.options = ok_gears.columns.levels[1]

    pyalgo_columns = [
        ("/".join(cc for cc in c if cc), c)
        for c in pyalgo.select_dtypes(exclude=[np.object])
        if not c[0].startswith("ok_gear") and not c[0] == "t"
    ]
    PyAlgoSignals.options = pyalgo_columns

    accdb_columns = [
        c for c in accdb.select_dtypes(exclude=[np.object]) if not c.startswith("Ind_g")
    ]
    AccDBSignals.options = accdb_columns

Case.observe(update_valid_signals, names="value")
update_valid_signals(None)

def apply_axis_scene(change):
    Case.value, Zoom.value, Pan.value, *other = AxisScenes.value
    Desc.value = other[0] if other else ""

AxisScenes.observe(apply_axis_scene, names="value")

def _count_series_diffs(a, b):
    na, nb = len(a), len(b)
    if na != nb:
        na = min(na, nb)
        return (a.iloc[:na] != b.iloc[:na]).sum() + (nb - na)
    return (a != b).sum()

def _distribute_signals_in_axes(columns, df):
    cols1, cols2 = decide_signal_axis(columns)
    df1 = df.loc[:, cols1]
    df2 = df.loc[:, cols2]
    return df1, df2

def recreate_fig():
    """
     Recreate the same figure-number, or else they leak.

    Hack, or else, shown double figure the 1st time this cell runs,
    or figure hidden, or stability/performance problems.
    """
    for i in range(3):
        fig_nums = plt.get_fignums()
        fig = plt.figure(
            num='Compare AccDB <--> PyAlgo "Initial Gear"', figsize=(10, 8)
        )
        if fig.number in set(fig_nums):
            fig.clear()
            plt.close(fig)
        else:
            return fig
    else:
        raise Exception(f"Exhausted new-fig tries({i}) with figs({plt.get_fignums()})")

fig = recreate_fig()

def plot_gear_flags(
    case,
    gear,
    zoom,
    pan,
    pyalgo_signals,
    accdb_signals,
    is_pyalgo_gears,
    is_accdb_gears,
):
    fig.clear()

    ax = plt.subplot()
    ax2 = ax3 = None

    out_specs = refresh_tabs()

    (
        cycle,
        ok_flags,
        ok_gears,
        accdb_cycle,
        accdb_gears,
        accdb_props,
    ) = load_interactive_case(case)

    clen = max(len(cycle), len(accdb_cycle))
    viewlen = int(clen / zoom)
    offset = int(pan * (clen - viewlen) / max_zoom)
    scale = idx[offset : offset + viewlen]

    diffgears_total = _count_series_diffs(cycle["g_max0"], accdb_cycle["g_max"])

    pyalgo = cycle.loc[scale]
    accdb = accdb_cycle.loc[scale]
    diffgears_view = _count_series_diffs(pyalgo["g_max0"], accdb["g_max"])
    merged = merge_pyalgo_accdb(pyalgo, accdb)

    ## Clean tabs if caseno/pan/zoom have changed.
    #
    if not pyalgo.equals(results[0]):
        Tab.selected_index = None
        for out, _, _ in out_specs:
            with out:
                clear_output()

    results[0] = pyalgo
    results[1] = accdb
    results[2] = accdb_props

    ax.set_title(
        f"Case_no: {case}, # of diff gears, in_view: {diffgears_view}, TOTAL: {diffgears_total}"
    )

    ok_flags = ok_flags.loc[:, idx[:, gear]].iloc[scale]
    ok_gears = ok_gears.loc[scale]

    if not pyalgo["V_cycle"].empty:
        pyalgo["V_cycle"].plot.line(ax=ax, color="0.70", linewidth=4.5)

    if is_pyalgo_gears:
        if not ok_flags.empty:
            ok_flags.plot.line(ax=ax, linewidth=2)
        ok_gears.plot.line(ax=ax, linewidth=3)

    ax2 = merged[["pyalgo/g_max0", "accdb/g_max"]].plot.line(
        ax=ax, linewidth=3, style=["b-", "c:"], secondary_y=True
    )

    if is_accdb_gears:
        accdb_gears = accdb.loc[:, accdb_gears] * (np.arange(len(accdb_gears)) + 1)
        accdb_gears[accdb_gears == 0] = np.NAN
        accdb_gears = accdb_gears.dropna(how="all", axis=1).fillna(0)
        ax2 = accdb_gears.plot.line(
            ax=ax, linewidth=2, linestyle="--", secondary_y=True
        )

    ax3 = ax.twinx()
    show_signals = pyalgo_signals or accdb_signals
    ax3.set_visible(show_signals)
    if show_signals:
        pyalgo1, pyalgo2 = _distribute_signals_in_axes(pyalgo_signals, pyalgo)
        accdb1, accdb2 = _distribute_signals_in_axes(accdb_signals, accdb)

        merged1 = merge_pyalgo_accdb(pyalgo1, accdb1)
        merged2 = merge_pyalgo_accdb(pyalgo2, accdb2)

        if not merged1.empty:
            merged1.astype(np.float64).plot.line(ax=ax, linewidth=2, linestyle="--")
        if not merged2.empty:
            merged2.astype(np.float64).plot.line(ax=ax3, linewidth=2, linestyle=":")
            ax3.legend(loc=4)
    ax.legend(loc=1)
    ax2.legend(loc=3)

    #     if ax2:
    #         ax2.grid(True, axis="both", which="both")
    ax.grid(True, axis="both", which="both")

    # Re-tighten if ax3 has (dis)appeared.
    # Had to use `rect` or axis-title half-hidden the first time fig is created!
    fig.tight_layout(rect=[0, 0, 1, 0.985])

display(Gui)

interactive_output(
    plot_gear_flags,
    {
        "case": Case,
        "gear": Gear,
        "zoom": Zoom,
        "pan": Pan,
        "pyalgo_signals": PyAlgoSignals,
        "accdb_signals": AccDBSignals,
        "is_pyalgo_gears": IsPyAlgoGears,
        "is_accdb_gears": IsAccdbGears,
    },
)

# %% [markdown]
# # Museum

# %%
## Is clutch-undefined only used for gear 1?  NO< also for g2.
display(c1.loc[c1.clutch == "undefined", "gear"].value_counts())
# display(c1.loc[c1.clutch=='undefined', ['v', 'a', 'g_max', 'clutch']])

# %%
# c2 = c2.loc[c2.index != 'v116']
# p2 = p2.loc[p2.index != 'v116']
# p2.to_excel('pyalgo_props-1.0.0.dev10.xlsx')
# c2.to_excel('pyalgo_cycles-1.0.0.dev10.xlsx')

# %% [markdown]
# ### AccDB not respecting n_min=0.9 x n_idle (Annex 2-3.k.3):

# %%
def is_bad_g2_in_decel_to_stop(accdb, cyc, prop):
    cyc = cyc.reset_index(level=0)
    # bad_rows = (accdb.g_max == 2) & cyc.stopdecel & (cyc['n/g2'] < 0.9 * prop['idling_speed'])
    bad_rows = (accdb.g_max == 2) & (cyc.g_max0 == 1) & cyc.stopdecel
    if bad_rows.any():
        return pd.concat(
            (
                accdb.loc[bad_rows, ["g_max", "n_2"]],
                cyc.loc[bad_rows, ["n/g2", "ok_gear/g1", "ok_gear/g2", "stopdecel"]],
            ),
            axis=1,
        )

for case, cyc in c2.groupby(level=0):
    accdb = c1.loc[case]
    prop = p1.loc[case]
    res = is_bad_g2_in_decel_to_stop(accdb, cyc, prop)
    if res is not None:
        display(
            prop[["idling_speed", "vehicle_class"]],
            f"0.9 x n_idle: {prop['idling_speed'] * 0.9}",
        )
        if isinstance(res, tuple):
            display(*res)
        else:
            display(res)

# %% [markdown]
# ### Insufficient power where more than one gears are N-valid:

# %%
def is_more_low_powered_gears(cyc):
    cyc2 = cyc.copy()
    cyc2.columns = wio.inflate_columns(cyc2.columns)
    c22 = cyc2.iloc[1571:1579][["ok_p", "ok_max_n"]].dropna(axis=1, how="all")
    try:
        ## is there any row with all-low-p AND 2-or-more n-max-ok?
        bad_rows = (~c22["ok_p"].replace(-1, 0).astype("bool")).all(axis=1) & (
            c22["ok_max_n"].replace(-1, 0).astype("bool").sum(axis=1) > 1
        )
        if bad_rows.any():
            c22.columns = wio.flatten_columns(c22.columns)
            return (
                list(bad_rows[bad_rows].index),
                pd.concat(
                    (c22, cyc.iloc[1571:1579].loc[:, ["g_min", "g_max"]]), axis=1
                ),
            )
    except Exception as ex:
        print(ex)

for case, cyc in c2.groupby(level=0):
    res = is_more_low_powered_gears(cyc)
    if res is not None:
        display(p2.loc[case, ["wltc_class", "g_vmax"]])
        if isinstance(res, tuple):
            display(*res)
        else:
            display(res)


# ---
# jupyter:
#   jupytext:
#     cell_metadata_filter: all
#     notebook_metadata_filter: all,-language_info,-toc,-latex_envs
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.6.0
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %%
import a301_lib
from pathlib import Path
from matplotlib import pyplot as plt
import pprint
import geopandas as gpd
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import cartopy
from pathlib import Path
import pprint
import numpy as np
from pyproj import CRS, Transformer
import datetime
import pytz
from IPython.display import display
pacific = pytz.timezone("US/Pacific")
date = datetime.datetime.today().astimezone(pacific)
print(f"written on {date}")

# %% [markdown]
# # Adding features to a cartopy map
#
# This is an updated version of the [first mapping notebook](https://a301_web.eoas.ubc.ca/week4/cartopy_mapping_h5.html#geographic-coordinate-systems) with two
# changes:
#
# 1. Change the coordinate transformation code from cartopy's [transform_points](https://scitools.org.uk/cartopy/docs/latest/crs/index.html) to pyproj's [Transformer](https://pyproj4.github.io/pyproj/stable/api/transformer.html)
#
# 2. Add features from a geojson file that maps North American rivers, which I downloaded from [natural earth](https://github.com/nvkelso/natural-earth-vector/tree/master/geojson)
#
# Why am I changing the code for CRS transformation from cartopy to pyproj?  Basically because cartopy is planning
# to move from its own transformation code to the newer pyproj format, but it's
# [still in underway](https://github.com/SciTools/cartopy/pull/1023).  Eventually, we won't have to switch between
# cartopy's coordinate objects and pyprojs, but for now, I need to create separate versions for each package:
#
#
#    1. Lambert Azimuthal Equal Area for cartopy: cartopy_laea
#    1. Lambert Azimuthal Equal Area for pyproj:  proj_laea
#    1. Geodetic lat/lon for cartopy: cartopy_latlon (cartopy.crs.PlateCarree())
#    1. Geodetic lat/lon for pyproj:  proj_latlon
#
# Note the different formats when I print them out below -- pyproj is much fancier.

# %% scrolled=false
#
cartopy_laea = ccrs.LambertAzimuthalEqualArea(
    central_latitude= 45, central_longitude=-123
)

proj_laea = CRS.from_proj4(cartopy_laea.proj4_init)
proj_latlon = CRS.from_proj4("+proj=latlon")
cartopy_latlon = cartopy.crs.PlateCarree()
print(f"{cartopy_latlon.proj4_params=}\n")
print(f"{cartopy_laea.proj4_params=}\n")
print(f"{proj_latlon=}")
print(f"{proj_laea=}")
print(f"{proj_latlon.to_wkt()=}")

# %% [markdown]
# ## Checking the coordinates
#
# In this cell I set up the bounding box.  As a santity check, I make sure that
# the upper left and lower right corner coordinates are in the correct order
# (left more negative than right, bottom more negative than top) and that
# Vancouver is inside the box.

# %%
ul_corner = (-135,52)
lr_corner = (-105,35)
transform = Transformer.from_crs(proj_latlon, proj_laea)
laea_x, laea_y = transform.transform([ul_corner[0],lr_corner[0]],
                                           [ul_corner[1],lr_corner[1]])
ul_corner = laea_x[0],laea_y[0]
lr_corner = laea_x[1], laea_y[1]
print(f"{[ul_corner,lr_corner]}=")
van_lon, van_lat = [-123.1207, 49.2827]
van_x, van_y = transform.transform(van_lon, van_lat)
print(f"{van_x=},{van_y=}")

# %%
fig, ax = plt.subplots(1, 1, figsize=(15, 15), subplot_kw={"projection": cartopy_laea})
#
# extent order  [xleft, xright, ybot, ytop]
#
laea_extent = [ul_corner[0], lr_corner[0], lr_corner[1], ul_corner[1]]
ax.set_extent(laea_extent, cartopy_laea)
#
# the simple lon,lat projection is called "geodetic"
#
ax.plot(van_x, van_y, "ro", markersize=10)
ax.gridlines(linewidth=2)
ax.add_feature(cartopy.feature.GSHHSFeature(scale="coarse", levels=[1, 2, 3]));
ax.coastlines(resolution="10m", color="red", lw=2);

# %% [markdown]
# ## Adding a new set of features
#
# Here are the shape files I've downloaded.  See the Readme_gshhs_wdbII.md for details.  The
# "10m" in the file names mean:  1:10 million -- i.e. 1 meter on a map is 10 million meters
# in the real world, or equivalently, 1 cm on the map is 100 km in the real world.

# %% scrolled=false
openstreetmap_dir = a301_lib.data_share / 'openstreetmap'
all_shapes = list(openstreetmap_dir.glob("*"))
[print(item.name) for item in all_shapes];

# %% [markdown]
# ## Read the North American Rivers geojson file

# %% scrolled=true
map_folder = a301_lib.data_share / 'openstreetmap'
na_rivers = list(map_folder.glob("*rivers*north*"))[0]
df_rivers=gpd.read_file(na_rivers)
print(f"{df_rivers.crs=}\n")
print(f"\n{df_rivers.head()=}\n")

# %% [markdown]
# * Here is the first row in the dataframe

# %%
df_rivers.iloc[0].geometry

# %% [markdown]
# ## Add rivers to the map
#
# Now put the rivers on the map and redraw.

# %% scrolled=false
ax.add_geometries(df_rivers['geometry'],cartopy_latlon,facecolor="none",edgecolor="green")
display(fig)


r""".. role:: html(raw)
   :format: html

Understanding the Haar Measure
==============================

.. meta::
    :property="og:description": Learn all about the Haar measure and how to randomly sample quantum states.

    :property="og:image": https://pennylane.ai/qml/_images/spherical_int_dtheta.png

.. related::

    tutorial_unitary_designs Unitary designs
    quantum_volume Quantum volume
    qsim_beyond_classical Beyond classical computing with qsim
    tutorial_barren_plateaus Barren plateaus

*Author: PennyLane dev team. Posted: 22 March 2021. Last updated: 22 March 2021.*

If you've ever dug into the literature about random quantum circuits,
variational ansatz structure, or anything related to the structure and
properties of unitary operations, you've likely come across a statement like the
following: "Assume that :math:`U` is sampled uniformly at random from the Haar
measure".  In this demo, we're going to unravel this cryptic statement and take
an in-depth look at what it means. You'll gain an understanding of the general
concept of *measure*, the Haar measure and its special properties, and you'll
learn how to sample from it using tools available in PennyLane and other
scientific computing frameworks. By the end of this demo, you'll be able to
include that important statement in your own work with confidence!

.. note::

   To get the most out of this demo, it is helpful if you are familiar with
   `integration of multi-dimensional functions
   <https://en.wikipedia.org/wiki/Multiple_integral>`__, the `Bloch sphere
   <https://en.wikipedia.org/wiki/Bloch_sphere>`__, and the conceptual ideas
   behind `decompositions
   <https://en.wikipedia.org/wiki/Matrix_decomposition>`__ and factorizations of
   unitary matrices (see, e.g., 4.5.1 and 4.5.2 of [#NandC2000]_).

Measure
-------

`Measure theory <https://en.wikipedia.org/wiki/Measure_(mathematics)>`__ is a
branch of mathematics that studies things that are measurable---think length,
area, or volume, but generalized to mathematical spaces and even higher
dimensions. Loosely, the measure tells you about how "stuff" is distributed and
concentrated in a mathematical set or space. An intuitive way to understand
measure is to think about a sphere. An arbitrary point on a sphere can be
parametrized by three numbers---depending on what you're doing, you may use
Cartesian coordinates :math:`(x, y, z)`, or it may be more convenient to use
spherical coordinates :math:`(\rho, \phi, \theta)`.

Suppose you wanted to compute the volume of a solid sphere with radius
:math:`r`.  This can be done by integrating over the three coordinates
:math:`\rho, \phi`, and :math:`\theta`. Your first thought here may be to simply
integrate each parameter over its full range, like so:

.. math::

    V = \int_0^{r} \int_0^{2\pi} \int_0^{\pi} d\rho~ d\phi~ d\theta = 2\pi^2 r

But, we know that the volume of a sphere of radius :math:`r` is
:math:`\frac{4}{3}\pi r^3`, so what we got from this integral is clearly wrong!
Taking the integral naively like this doesn't take into account the structure of
the sphere with respect to the parameters. For example, consider
two small, infinitesimal elements of area with the same difference in
:math:`\theta` and :math:`\phi`, but at different values of :math:`\theta`:

.. figure:: /demonstrations/haar_measure/spherical_int_dtheta.png
    :align: center
    :width: 50%

    |

Even though the differences :math:`d\theta` and :math:`d\phi` themselves are the
same, there is way more "stuff" near the equator of the sphere than there is
near the poles. We must take into account the value of :math:`\theta` when
computing the integral! Specifically, we multiply by the function
:math:`\sin\theta`---the properties of the :math:`\sin` function mean that the
most weight will occur around the equator where :math:`\theta=\pi/2`, and the
least weight near the poles where :math:`\theta=0` and :math:`\theta=\pi`.

Similar care must be taken for :math:`\rho`.  The contribution to volume of
parts of the sphere with a large :math:`\rho` is far more than for a small
:math:`\rho`---we should expect the contribution to be proportional to
:math:`\rho^2`, given that the surface area of a sphere of radius :math:`r` is
:math:`4\pi r^2`.

On the other hand, for a fixed :math:`\rho` and :math:`\theta`, the length of
the :math:`d\phi` is the same all around the circle. If put all these facts
together, we find that the actual expression for the integral should look like
this:

.. math::

    V = \int_0^r \int_0^{2\pi} \int_0^{\pi} \rho^2 \sin \theta~ d\rho~ d\phi~
    d\theta = \frac{4}{3}\pi r^3

These extra terms that we had to add to the integral, :math:`\rho^2 \sin
\theta`, constitute the *measure*. The measure weights portions of the sphere
differently depending on where they are in the space. While we need to know the
measure to properly integrate over the sphere, knowledge of the measure also
gives us the means to perform another important task, that of sampling points in
the space uniformly at random. We can't simply sample each parameter from the
uniform distribution over its domain---as we experienced already, this doesn't
take into account how the sphere is spread out over space. The measure describes
the distribution of each parameter and gives a recipe for sampling them in order
to obtain something properly uniform.

The Haar measure
----------------

Operations in quantum computing are described by unitary matrices.
Unitary matrices, like points on a sphere, can be expressed in terms of a fixed
set of coordinates, or parameters. For example, the most general single-qubit rotation
implemented in PennyLane (:class:`~.pennylane.Rot`) is expressed in terms of three
parameters like so,

.. math::

    U(\phi, \theta, \omega) = \begin{pmatrix} e^{-i(\phi + \omega)/2}
                        \cos(\theta/2) & -e^{i(\phi - \omega)/2} \sin(\theta/2)
                        \\ e^{-i(\phi - \omega)/2} \sin(\theta/2) & e^{i(\phi +
                        \omega)/2} \cos(\theta/2) \end{pmatrix}.

For every dimension :math:`N`, the unitary matrices of size :math:`N \times N`
constitute the *unitary group* :math:`U(N)`. We can perform operations on
elements of this group, such as apply functions to them, integrate over them, or
sample uniformly over them, just as we can do to points on a sphere. When we do
such tasks with respect to the sphere, we have to add the measure in order to
properly weight the different regions of space. The *Haar measure* provides the
analogous terms we need for working with the unitary group.

For an :math:`N`-dimensional system, the Haar measure, often denoted by
:math:`\mu_N`, tells us how to weight the elements of :math:`U(N)`. For
example, suppose :math:`f` is a function that acts on elements of :math:`U(N)`,
and we would like to take its integral over the group. We must write this
integral with respect to the Haar measure, like so:

.. math::

    \int_{V \in U(N)} f(V) d\mu_N(V).

As with the measure term of the sphere, :math:`d\mu_N` itself can be broken down
into components depending on individual parameters.  While the Haar
measure can be defined for every dimension :math:`N`, the mathematical form gets
quite hairy for larger dimensions---in general, an :math:`N`-dimensional unitary
requires at least :math:`N^2 - 1` parameters, which is a lot to keep track of!
Therefore we'll start with the case of a single qubit :math:`(N=2)`, then show
how things generalize.

Single-qubit Haar measure
~~~~~~~~~~~~~~~~~~~~~~~~~

The single-qubit case provides a particularly nice entry point because we can
continue our comparison to spheres by visualizing single-qubit states on the
Bloch sphere. As expressed above, the measure provides a recipe for sampling
elements of the unitary group in a properly uniform manner, given the structure
of the group. One useful consequence of this is that it provides a method to
sample quantum *states* uniformly at random---we simply generate Haar-random
unitaries, and apply them to a fixed basis state such as :math:`\vert 0\rangle`.

We'll see how this works in good time. First, we'll take a look at what happens
when we ignore the measure and do things *wrong*. Suppose we sample quantum
states by applying unitaries obtained by the parametrization above, but sample
the angles :math:`\omega, \phi`, and :math:`\theta` from the flat uniform
distribution between :math:`[0, 2\pi)` (fun fact: there is a measure implicit in
this kind of sampling too! It just has a constant value, because each point is
equally likely to be sampled).

"""

import pennylane as qml
import numpy as np
import matplotlib.pyplot as plt

# set the random seed
np.random.seed(42)

# Use the mixed state simulator to save some steps in plotting later
dev = qml.device('default.mixed', wires=1)

@qml.qnode(dev)
def not_a_haar_random_unitary():
    # Sample all parameters from their flat uniform distribution
    phi, theta, omega = 2 * np.pi * np.random.uniform(size=3)
    qml.Rot(phi, theta, omega, wires=0)
    return qml.state()

num_samples = 2021

not_haar_samples = [not_a_haar_random_unitary() for _ in range(num_samples)]

######################################################################
# In order to plot these on the Bloch sphere, we'll need to do one more
# step, and convert the quantum states into Bloch vectors.
#

X = np.array([[0, 1], [1, 0]])
Y = np.array([[0, -1j], [1j, 0]])
Z = np.array([[1, 0], [0, -1]])

# Used the mixed state simulator so we could have the density matrix for this part!
def convert_to_bloch_vector(rho):
    """Convert a density matrix to a Bloch vector."""
    ax = np.trace(np.dot(rho, X)).real
    ay = np.trace(np.dot(rho, Y)).real
    az = np.trace(np.dot(rho, Z)).real
    return [ax, ay, az]

not_haar_bloch_vectors = np.array([convert_to_bloch_vector(s) for s in not_haar_samples])

######################################################################
# With this done, let's find out where our "uniformly random" states ended up:

def plot_bloch_sphere(bloch_vectors):
    """ Helper function to plot vectors on a sphere."""
    fig = plt.figure(figsize=(6, 6))
    ax = fig.add_subplot(111, projection='3d')
    fig.subplots_adjust(left=0, right=1, bottom=0, top=1)

    ax.grid(False)
    ax.set_axis_off()
    ax.view_init(30, 45)
    ax.dist = 7

    # Draw the axes (source: https://github.com/matplotlib/matplotlib/issues/13575)
    x, y, z = np.array([[-1.5,0,0], [0,-1.5,0], [0,0,-1.5]])
    u, v, w = np.array([[3,0,0], [0,3,0], [0,0,3]])
    ax.quiver(x, y, z, u, v, w, arrow_length_ratio=0.05, color="black", linewidth=0.5)

    ax.text(0, 0, 1.7, r"|0⟩", color="black", fontsize=16)
    ax.text(0, 0, -1.9, r"|1⟩", color="black", fontsize=16)
    ax.text(1.9, 0, 0, r"|+⟩", color="black", fontsize=16)
    ax.text(-1.7, 0, 0, r"|–⟩", color="black", fontsize=16)
    ax.text(0, 1.7, 0, r"|i+⟩", color="black", fontsize=16)
    ax.text(0,-1.9, 0, r"|i–⟩", color="black", fontsize=16)

    ax.scatter(
        bloch_vectors[:,0], bloch_vectors[:,1], bloch_vectors[:, 2], c='#e29d9e', alpha=0.3
    )

plot_bloch_sphere(not_haar_bloch_vectors)

######################################################################
# You can see from this plot that even though our parameters were sampled from a
# uniform distribution, there is a noticeable amount of clustering around the poles
# of the sphere. Despite the input parameters being uniform, the output is very
# much *not* uniform. Just like the regular sphere, the measure is larger near
# the equator, and if we just sample uniformly, we won't end up populating that
# area as much. To take that into account we will need to sample from the proper
# Haar measure, and weight the different parameters appropriately.
#
# For a single qubit, the Haar measure looks much like the case of a sphere,
# minus the radial component. Intuitively, all qubit state vectors have length
# 1, so it makes sense that this wouldn't play a role here. The parameter that
# we will have to weight differently is :math:`\theta`, and in fact the
# adjustment in measure is identical to that we had to do with the polar axis of
# the sphere, i.e., :math:`\sin \theta`. In order to sample the :math:`\theta`
# uniformly at random in this context, we must sample from the distribution
# :math:`\hbox{Pr}(\theta) = \sin \theta`. We can accomplish this by setting up
# a custom probability distribution with
# `rv_continuous <https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.rv_continuous.html#scipy.stats.rv_continuous>`__
# in ``scipy``.

from scipy.stats import rv_continuous

class sin_prob_dist(rv_continuous):
    def _pdf(self, theta):
        # The 0.5 is so that the distribution is normalized
        return 0.5 * np.sin(theta)

# Samples of theta should be drawn from between 0 and pi
sin_sampler = sin_prob_dist(a=0, b=np.pi)

@qml.qnode(dev)
def haar_random_unitary():
    phi, omega = 2 * np.pi * np.random.uniform(size=2) # Sample phi and omega as normal
    theta = sin_sampler.rvs(size=1) # Sample theta from our new distribution
    qml.Rot(phi, theta, omega, wires=0)
    return qml.state()

haar_samples = [haar_random_unitary() for _ in range(num_samples)]
haar_bloch_vectors = np.array([convert_to_bloch_vector(s) for s in haar_samples])

plot_bloch_sphere(haar_bloch_vectors)

######################################################################
# We see that when we use the correct measure, our qubit states are now
# much better distributed over the sphere. Putting this information together,
# we can now write the explicit form for the single-qubit Haar measure:
#
# .. math::
#
#    d\mu_2 = \sin \theta d\theta \cdot d\omega \cdot d\phi.
#
# Show me more math!
# ~~~~~~~~~~~~~~~~~~
#
# While we can easily visualize the single-qubit case, this is no longer
# possible when we increase the number of qubits. Regardless, we can still
# obtain a mathematical expression for the Haar measure in arbitrary
# dimensions. In the previous section, we expressed the Haar measure in terms of
# a set of parameters that can be used to specify the unitary group
# :math:`U(2)`. Such a parametrization is not unique, and in fact there are
# multiple ways to *factorize*, or decompose an :math:`N`-dimensional unitary
# operation into a set of parameters.
#
# Many of these parametrizations come to us from the study of photonics.  Here,
# arbitrary operations are broken down into elementary operations involving only
# a few parameters which correspond directly to parameters of the physical
# apparatus used to implement them (beamsplitters and phase shifts). Rather than
# qubits, such operations act on modes, or *qumodes*. They are expressed as
# elements of the :math:`N`-dimensional `special unitary group
# <https://en.wikipedia.org/wiki/Special_unitary_group>`__. This group, written
# as :math:`SU(N)`, is the continuous group consisting of all :math:`N \times N`
# unitary operations with determinant 1 (essentially like :math:`U(N)`, minus
# a potential global phase).
#
#
# .. note::
#
#     Elements of :math:`SU(N)` and :math:`U(N)` can still be considered as
#     multi-qubit operations in the cases where :math:`N` is a power of 2, but
#     they must be translated from continuous-variable operations into qubit
#     operations. (In PennyLane, this can be done by feeding the unitaries to
#     the :class:`~.pennylane.QubitUnitary` operation directly. Alternatively,
#     one can use *quantum compilation* to express the operations as a sequence
#     of elementary gates such as Pauli rotations and CNOTs.)
#
# .. admonition:: Tip
#
#    If you haven't had many opportunities to work in terms of qumodes, the
#    `Strawberry Fields documentation
#    <https://strawberryfields.ai/photonics/concepts/photonics.html>`__ is a
#    good starting point.
#
# For example, we saw already above that for :math:`N=2`, we can write
#
# .. math::
#
#    U(\phi, \theta, \omega) = \begin{pmatrix} e^{-i(\phi + \omega)/2}
#                        \cos(\theta/2) & -e^{i(\phi - \omega)/2} \sin(\theta/2)
#                        \\ e^{-i(\phi - \omega)/2} \sin(\theta/2) & e^{i(\phi +
#                        \omega)/2} \cos(\theta/2) \end{pmatrix}.
#
#
# This unitary can be factorized as follows:
#
# .. math::
#
#    U(\phi, \theta, \omega) =
#        \begin{pmatrix}
#          e^{-i\omega/2} & 0 \\ 0 & e^{i\omega/2}
#        \end{pmatrix}
#        \begin{pmatrix}
#          \cos(\theta/2) & -\sin(\theta/2) \\ \sin(\theta/2) & \cos(\theta/2)
#        \end{pmatrix}
#       \begin{pmatrix}
#          e^{-i\phi/2} & 0 \\ 0 & e^{i\phi/2}
#        \end{pmatrix}
#
# The middle operation is a beamsplitter; the other two operations are phase
# shifts.  We saw earlier that for :math:`N=2`, :math:`d\mu_2 = \sin\theta
# d\theta d\omega d\phi`---note how the parameter in the beamsplitter
# contributes to the measure in a different way than those of the phase
# shifts. As mentioned above, for larger values of :math:`N` there are multiple
# ways to decompose the unitary. Such decompositions rewrite elements in
# :math:`SU(N)` acting on :math:`N` modes as a sequence of operations acting
# only on 2 modes, :math:`SU(2)`, and single-mode phase shifts.  Shown below are
# three examples [#deGuise2018]_, [#Clements2016]_, [#Reck1994]_:
#
# .. figure:: /demonstrations/haar_measure/unitaries.png
#    :align: center
#    :width: 95%
#
#
# In these graphics, every wire is a different mode. Every box represents an
# operation on one or more modes, and the number in the box indicates the number
# of parameters.  The boxes containing a ``1`` are simply phase shifts on
# individual modes. The blocks containing a ``3`` are :math:`SU(2)` transforms
# with 3 parameters, such as the :math:`U(\phi, \theta, \omega)` above. Those
# containing a ``2`` are :math:`SU(2)` transforms on pairs of modes with 2
# parameters, similar to the 3-parameter ones but with :math:`\omega = \phi`.
#
# Although the decompositions all produce the same set of operations, their
# structure and parametrization may have consequences in practice.  The first [#deGuise2018]_
# has a particularly convenient form that leads to a recursive definition
# of the Haar measure. The decomposition is formulated recursively such that an
# :math:`SU(N)` operation can be implemented by sandwiching an :math:`SU(2)`
# transformation between two :math:`SU(N-1)` transformations, like so:
#
# |
#
# .. figure:: /demonstrations/haar_measure/sun.svg
#    :align: center
#    :width: 80%
#
# |
#
# The Haar measure is then constructed recursively as a product of 3
# terms. The first term depends on the parameters in the first :math:`SU(N-1)`
# transformation; the second depends on the parameters in the lone :math:`SU(2)`
# transformation; and the third term depends on the parameters in the other
# :math:`SU(N-1)` transformation.
#
# :math:`SU(2)` is the "base case" of the recursion---we simply have the Haar measure
# as expressed above.
#
# |
#
# .. figure:: /demonstrations/haar_measure/su2_haar.svg
#    :align: center
#    :width: 25%
#
# |
#
# Moving on up, we can write elements of :math:`SU(3)` as a sequence of three
# :math:`SU(2)` transformations. The Haar measure :math:`d\mu_3` then consists
# of two copies of :math:`d\mu_2`, with an extra term in between to take into
# account the middle transformation.
#
# |
#
# .. figure:: /demonstrations/haar_measure/su3_haar.svg
#    :align: center
#    :width: 80%
#
# |
#
# For :math:`SU(4)` and upwards, the form changes slightly, but still follows
# the pattern of two copies of :math:`d\mu_{N-1}` with a term in between.
#
# |
#
# .. figure:: /demonstrations/haar_measure/su4_premerge.svg
#    :align: center
#    :width: 90%
#
# |
#
# For larger systems, however, the recursive composition allows for some of the
# :math:`SU(2)` transformations on the lower modes to be grouped. We can take
# advantage of this and aggregate some of the parameters:
#
# |
#
# .. figure:: /demonstrations/haar_measure/su4_triangle_merge.svg
#    :align: center
#    :width: 100%
#
# |
#
# This leads to one copy of :math:`d\mu_{N-1}`, which we'll denote as
# :math:`d\mu_{N-1}^\prime`, containing only a portion of the full set of terms
# (as detailed in [#deGuise2018]_, this is called a *coset measure*).
#
# |
#
# .. figure:: /demonstrations/haar_measure/su4_haar.svg
#    :align: center
#    :width: 100%
#
# |
#
# Putting everything together, we have that
#
# .. math::
#
#    d\mu_N = d\mu_{N-1}^\prime \times \sin \theta_{N-1}
#    \sin^{2(N-2)}\left(\frac{\theta_{N-1}}{2}\right) d\theta_{N-1} d\omega_{N-1} \times d\mu_{N-1}
#
# The middle portion depends on the value of :math:`N`, and the parameters
# :math:`\theta_{N-1}` and :math:`\omega_{N-1}` contained in the :math:`(N-1)`'th
# :math:`SU(N)` transformation. This is thus a convenient, systematic way to
# construct the :math:`N`-dimensional Haar measure for the unitary group. As a
# final note, even though unitaries can be parametrized in different ways, the
# underlying Haar measure is *unique*. This is a consequence of it being an
# invariant measure, as will be shown later.
#
# Haar-random matrices from the :math:`QR` decomposition
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# Nice-looking math aside, sometimes you just need to generate a large number of
# high-dimensional Haar-random matrices. It would be very cumbersome to sample
# and keep track of the distributions of so many parameters; furthermore, the
# measure above requires you to parametrize your operations in a fixed way.
# There is a much quicker way to perform the sampling by taking a (slightly
# modified) `QR decomposition
# <https://en.wikipedia.org/wiki/QR_decomposition>`__ of complex-valued
# matrices.  This algorithm is detailed in [#Mezzadri2006]_, and consists of the
# following steps:
#
# 1. Generate an :math:`N \times N` matrix :math:`Z` with complex numbers :math:`a+bi`
#    where both :math:`a` and :math:`b` are normally distributed with mean 0 and variance 1
#    (this is sampling from the distribution known as the *Ginibre ensemble*).
# 2. Compute a QR decomposition :math:`Z = QR`.
# 3. Compute the diagonal matrix :math:`\Lambda = \hbox{diag}(R_{ii}/|R_{ii}|)`.
# 4. Compute :math:`Q^\prime = Q \Lambda`, which will be Haar-random.
#
#

from numpy.linalg import qr

def qr_haar(N):
    """Generate a Haar-random matrix using the QR decomposition."""
    # Step 1
    A, B = np.random.normal(size=(N, N)), np.random.normal(size=(N, N))
    Z = A + 1j * B

    # Step 2
    Q, R = qr(Z)

    # Step 3
    Lambda = np.diag([R[i, i] / np.abs(R[i, i]) for i in range(N)])

    # Step 4
    return np.dot(Q, Lambda)

######################################################################
# Let's check that this method actually generates Haar-random unitaries
# by trying it out for :math:`N=2` and plotting on the Bloch sphere.
#

@qml.qnode(dev)
def qr_haar_random_unitary():
    qml.QubitUnitary(qr_haar(2), wires=0)
    return qml.state()

qr_haar_samples = [qr_haar_random_unitary() for _ in range(num_samples)]
qr_haar_bloch_vectors = np.array([convert_to_bloch_vector(s) for s in qr_haar_samples])
plot_bloch_sphere(qr_haar_bloch_vectors)

######################################################################
# As expected, we find our qubit states are distributed uniformly over the
# sphere.  This particular method is what's implemented in packages like
# ``scipy``'s `unitary_group
# <https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.unitary_group.html>`__
# function.
#
# Now, it's clear that this method works, but it is also important to
# understand *why* it works.  Step 1 is fairly straightforward---the base of our
# samples is a matrix full of complex values chosen from a typical
# distribution. This isn't enough by itself, since unitary matrices also
# have constraints---their rows and columns must be orthonormal.
# These constraints are where step 2 comes in---the outcome of a generic
# QR decomposition consists of an *orthonormal* matrix :math:`Q`, and and upper
# triangular matrix :math:`R`. Since our original matrix was complex-valued, we end
# up with a :math:`Q` that is in fact already unitary. But why not stop there? Why
# do we then perform steps 3 and 4?
#
# Steps 3 and 4 are needed because, while the QR decomposition yields a unitary,
# it is not a unitary that is properly Haar-random. In [#Mezzadri2006]_, it is
# explained that a uniform distribution over unitary matrices should also yield
# a uniform distribution over the *eigenvalues* of those matrices, i.e., every
# eigenvalue should be equally likely. Just using the QR decomposition out of
# the box produces an *uneven* distribution of eigenvalues of the unitaries!
# This discrepancy stems from the fact that the QR decomposition is not unique.
# We can take any unitary diagonal matrix :math:`\Lambda`, and re-express the decomposition
# as :math:`QR = Q\Lambda \Lambda^\dagger R = Q^\prime R^\prime`. Step 3 removes this
# redundancy by fixing a :math:`\Lambda` that depends on :math:`R`, leading to a unique
# value of :math:`Q^\prime = Q \Lambda`, and a uniform distribution of eigenvalues.
#
# .. admonition:: Try it!
#
#    Use the ``qr_haar`` function above to generate random unitaries and construct
#    a distribution of their eigenvalues. Then, comment out the lines for steps 3 and
#    4 and do the same---you'll find that the distribution is no longer uniform.
#    Check out reference [#Mezzadri2006]_ for additional details and examples.

######################################################################
# Fun (and not-so-fun) facts
# --------------------------
#
# We've now learned what the Haar measure is, and both an analytical and
# numerical means of sampling quantum states and unitary operations uniformly at
# random. The Haar measure also has many neat properties that play a role in
# quantum computing.
#
# Invariance of the Haar measure
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# Earlier, we showed that the Haar measure is used when integrating functions over
# the unitary group:
#
# .. math::
#
#    \int_{V \in U(N)} f(V) d\mu_N(V).
#
# One of the defining features of the Haar measure is that it is both left and
# right *invariant* under unitary transformations. That is,
#
# .. math::
#
#    \int_{V \in U(N)} f(\color{red}{W}V) d\mu_N(V) =  \int_{V \in U(N)} f(V\color{red}{W}) d\mu_N(V) =  \int_{V \in U(N)} f(V) d\mu_N(V).
#
# This holds true for *any* other :math:`N\times N` unitary :math:`W`! A
# consequence of such invariance is that if :math:`V` is Haar-random, then so is
# :math:`V^T,` :math:`V^\dagger,` and any product of another unitary matrix and
# :math:`V` (where the product may be taken on either side).
#
# Another consequence of this invariance has to do with the structure of the entries
# themselves: they must all come from the same distribution. This is because the
# measure remains invariant under permutations, since permutations are unitary---
# the whole thing still has to be Haar random no matter how the entries are ordered,
# so all distributions must be the same.  The specific distribution is complex
# numbers :math:`a+bi` where both :math:`a` and :math:`b` has mean 0 and variance
# :math:`1/N` [#Meckes2014]_ (so, much like Ginibre ensemble we used in the QR decomposition
# above, but with a different variance and constraints due to orthonormality).
#
# Concentration of measure
# ~~~~~~~~~~~~~~~~~~~~~~~~
#
# An unfortunate (although interesting) property of the Haar measure is that it
# suffers from the phenomenon of `concentration of measure
# <https://en.wikipedia.org/wiki/Concentration_of_measure>`__. Most of the
# "stuff" in the space concentrates around a certain area, and this gets worse
# as the size of the system increases. You can see the beginnings of by looking
# at the sphere. For the 3-dimensional sphere, we saw graphically how there is
# concentration around the equator, and how the measure takes that into account
# with the additional factor of :math:`\sin \theta`. This property becomes
# increasingly prominent for `higher-dimensional spheres
# <https://en.wikipedia.org/wiki/N-sphere>`__.
#
# .. important::
#
#    The concentration described here is not referring to what we witnessed
#    earlier on, when we sampled quantum states (points on the Bloch sphere)
#    incorrectly and found that they clustered around the poles. However, that
#    is not unrelated. Concentration of measure refers to where the measure
#    itself is concentrated, and which parts of the space should be more heavily
#    weighted. For the case of the sphere, it is the equatorial area, and when
#    we didn't sample properly and take that concentration into account, we
#    obtained an uneven distribution.
#
# Let's consider an :math:`N`-dimensional unit sphere. Points on the sphere, or
# vectors in this space, are parametrized by :math:`N-1` real coordinates.
# Suppose we have some function :math:`f` that maps points on that sphere to
# real numbers. Sample a point :math:`x` on that sphere from the uniform
# measure, and compute the value of :math:`f(x)`. How close do you think the
# result will be to the mean value of the function, :math:`E[f]`, over the
# entire sphere?
#
# A result called `Levy's lemma
# <https://en.wikipedia.org/wiki/Concentration_of_measure#Concentration_on_the_sphere>`__
# [#Gerken2013]_, [#Hayden2006]_ expresses how likely it is that :math:`f(x)` is a specific
# distance away from the mean. It states that, for an :math:`x` selected
# uniformly at random, the probability that :math:`f(x)` deviates from
# :math:`E[f]` by some amount :math:`\epsilon` is bounded by:
#
# .. math::
#
#    \hbox{Pr}(|f(x) - E[f]| \ge \epsilon) \leq 2 \exp\left[-\frac{N\epsilon^2}{9\pi^3 \eta^2}\right].
#
# A constraint on the function :math:`f` is that it must be `Lipschitz
# continuous <https://en.wikipedia.org/wiki/Lipschitz_continuity>`__, where
# :math:`\eta` is the *Lipschitz constant* of the function. The important aspect
# here is the likelihood of deviating significantly from the mean by an amount
# :math:`\epsilon` decreases exponentially with :math:`\epsilon.` Furthermore,
# increasing the dimension :math:`N` also makes the deviation exponentially less
# likely.
#
# Now, this result seems unrelated to quantum states---it concerns higher-
# dimensional spheres. However, recall that a quantum state vector is a complex
# vector whose squared values sum to 1, similar to vectors on a sphere. If you
# "unroll" a quantum state vector of dimension :math:`N = 2^n` by stacking its
# real and complex parts, you end with a vector of length :math:`2 \cdot 2^{n}`
# which ends up behaving just like a unit vector on the sphere in this
# dimension. Given that measure concentrates on spheres, and quantum state
# vectors can be converted to vectors on spheres, functions on random quantum
# states will also demonstrate concentration.
#
# This is bad news! To do useful things in quantum computing, we need a lot of
# qubits. But the more qubits we have, the more our randomly sampled states will
# look the same (specifically, random states will concentrate around the
# maximally entangled state [#Hayden2006]_). This has important consequences for
# near-term algorithms (as detailed in the next section), and any algorithm that
# involves uniform sampling of quantum states and operations.
#
# Haar measure and barren plateaus
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# Suppose you are venturing out to solve a new problem using an algorithm such
# as the :doc:`variational quantum eigensolver </demos/tutorial_vqe>`. A
# critical component of such methods is the choice of :doc:`variational ansatz
# </glossary/circuit_ansatz>`. Having now learned a bit about the properties of
# the Haar measure, you may think it would make sense to use this for the
# parametrization. Variational ansaetze are, after all, parametrized quantum
# circuits, so why not choose an ansatz that corresponds directly to a
# parametrization for Haar-random unitaries?  The initial parameter selection
# will give you a state in the Hilbert space uniformly at random. Then, since
# this ansatz spans the entire Hilbert space, you're guaranteed to be able to
# represent the target ground state with your ansatz, and it should be able to
# find it with no issue ... right?
#
# Unfortunately, while such an ansatz is extremely *expressive* (i.e., it is
# capable of representing any possible state), these ansaetze actually suffer
# the most from the barren plateau problem [#McClean2018]_, [#Holmes2021]_.
# :doc:`Barren plateaus </demos/tutorial_barren_plateaus>` are regions in the
# cost landscape of a parametrized circuit where both the gradient and its
# variance approach 0, leading the optimizer to get stuck in a local minimum.
# This was explored recently in the work of [#Holmes2021]_, wherein closeness to
# the Haar measure was actually used as a metric for expressivity. The closer
# things are to the Haar measure, the more expressive they are, but they are
# also more prone to exhibiting barren plateaus.
#
#
# .. figure:: /demonstrations/haar_measure/holmes-costlandscapes.png
#    :align: center
#    :width: 50%
#
#    Image source: [#Holmes2021]_. A highly expressive ansatz that can access
#    much of the space of possible unitaries (i.e., an ansatz capable of
#    producing unitaries in something close to a Haar-random manner) is very
#    likely to have flat cost landscapes and suffer from the barren plateau
#    problem.
#
# It turns out that the types of ansaetze know as *hardware-efficient ansaetze*
# also suffer from this problem if they are "random enough" (this notion will be
# formalized in a future demo). It was shown in [#McClean2018]_ that this is a
# consequence of the concentration of measure phenomenon described above. The
# values of gradients and variances can be computed for classes of circuits on
# average by integrating with respect to the Haar measure, and it is shown that
# these values decrease exponentially in the number of qubits, and thus huge
# swaths of the cost landscape are simply and fundamentally flat.
#
# Conclusion
# ----------
#
# The Haar measure plays an important role in quantum computing---anywhere
# you might be dealing with sampling random circuits, or averaging over
# all possible unitary operations, you'll want to do so with respect
# to the Haar measure.
#
# There are two important aspects of this that we have yet to touch upon,
# however. The first is whether it is efficient to sample from the Haar measure---given
# that the number of parameters to keep track of is exponential in the
# number of qubits, certainly not. But a more interesting question is do we
# *need* to always sample from the full Haar measure?  The answer to this is
# "no" in a very interesting way. Depending on the task at hand, you may be able
# to take a shortcut using something called a *unitary design*. In an upcoming
# demo, we will explore the amazing world of unitary designs and their
# applications!
#
# References
# ----------
#
# .. [#NandC2000]
#
#     M. A. Nielsen, and I. L. Chuang (2000) "Quantum Computation and Quantum Information",
#     Cambridge University Press.
#
# .. [#deGuise2018]
#
#     H. de Guise, O. Di Matteo, and L. L. Sánchez-Soto. (2018) "Simple factorization
#     of unitary transformations", `Phys. Rev. A 97 022328
#     <https://journals.aps.org/pra/abstract/10.1103/PhysRevA.97.022328>`__.
#     (`arXiv <https://arxiv.org/abs/1708.00735>`__)
#
# .. [#Clements2016]
#
#     W. R. Clements, P. C. Humphreys, B. J. Metcalf, W. S. Kolthammer, and
#     I. A. Walmsley (2016) “Optimal design for universal multiport
#     interferometers”, \ `Optica 3, 1460–1465
#     <https://www.osapublishing.org/optica/fulltext.cfm?uri=optica-3-12-1460&id=355743>`__.
#     (`arXiv <https://arxiv.org/abs/1603.08788>`__)
#
# .. [#Reck1994]
#
#    M. Reck, A. Zeilinger, H. J. Bernstein, and P. Bertani (1994) “Experimental
#    realization of any discrete unitary operator”, `Phys. Rev. Lett.73, 58–61
#    <https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.73.58>`__.
#
# .. [#Mezzadri2006]
#
#     F. Mezzadri (2006) "How to generate random matrices from the classical compact groups".
#     (`arXiv <https://arxiv.org/abs/math-ph/0609050>`__)
#
# .. [#Meckes2014]
#
#     E. Meckes (2019) `"The Random Matrix Theory of the Classical Compact Groups"
#     <https://case.edu/artsci/math/esmeckes/Haar_book.pdf>`_, Cambridge University Press.
#
# .. [#Gerken2013]
#
#     M. Gerken (2013) "Measure concentration: Levy's Lemma"
#     (`lecture notes <http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.679.2560>`__).
#
#
# .. [#Hayden2006]
#
#     P. Hayden, D. W. Leung, and A. Winter (2006) "Aspects of generic
#     entanglement", `Comm. Math. Phys. Vol. 265, No. 1, pp. 95-117
#     <https://link.springer.com/article/10.1007%2Fs00220-006-1535-6>`__.
#     (`arXiv <https://arxiv.org/abs/quant-ph/0407049>`__)
#
# .. [#McClean2018]
#
#     J. R. McClean, S. Boixo, V. N. Smelyanskiy, R. Babbush, and H. Neven
#     (2018) "Barren plateaus in quantum neural network training
#     landscapes", `Nature Communications, 9(1)
#     <http://dx.doi.org/10.1038/s41467-018-07090-4>`__.
#     (`arXiv <https://arxiv.org/abs/1803.11173>`__)
#
# .. [#Holmes2021]
#
#     Z. Holmes, K. Sharma, M. Cerezo, and P. J. Coles (2021) "Connecting ansatz
#     expressibility to gradient magnitudes and barren plateaus". (`arXiv
#     <https://arxiv.org/abs/2101.02138>`__)
#
#


import pickle
from numpy.core.numeric import isscalar
import pytest
from copy import deepcopy
import sys
from pathlib import Path
import numpy as np
import os
from dataclasses import is_dataclass, astuple
from collections.abc import Iterable

assignment_name = "eskf"

this_file = Path(__file__)
tests_folder = this_file.parent
test_data_file = tests_folder.joinpath("test_data.pickle")
project_folder = tests_folder.parent
code_folder = project_folder.joinpath(assignment_name)

sys.path.insert(0, str(code_folder))

import solution  # nopep8
import cross_matrix, eskf, nis_nees, quaternion  # nopep8

@pytest.fixture
def test_data():
    with open(test_data_file, "rb") as file:
        test_data = pickle.load(file)
    return test_data

def compare(a, b):
    if isinstance(b, np.ndarray) or np.isscalar(b):
        return np.allclose(a, b, atol=1e-6)

    elif is_dataclass(b):
        if type(a).__name__ != type(b).__name__:
            return False
        a_tup, b_tup = astuple(a), astuple(b)
        return all([compare(i, j) for i, j in zip(a_tup, b_tup)])

    elif isinstance(b, Iterable):
        return all([compare(i, j) for i, j in zip(a, b)])

    else:
        return a == b

class Test_get_cross_matrix:
    def test_output(self, test_data):
        """Tests if the function is correct by comparing the output
        with the output of the solution

        As python always use pass by reference, not by copy, it also checks if the
        input is changed (or not) in the same way as the in solution
        """
        for finput in test_data["cross_matrix.get_cross_matrix"]:
            params = tuple(finput.values())

            vec_1, = deepcopy(params)

            vec_2, = deepcopy(params)

            S_1 = cross_matrix.get_cross_matrix(vec_1,)

            S_2 = solution.cross_matrix.get_cross_matrix(vec_2,)

            assert compare(S_1, S_2)

            assert compare(vec_1, vec_2)

    def test_solution_usage(self, test_data):
        """Tests if the solution is used in the function"""
        for finput in test_data["cross_matrix.get_cross_matrix"][:1]:
            params = finput

            solution.used["cross_matrix.get_cross_matrix"] = False

            cross_matrix.get_cross_matrix(**params)

            assert not solution.used["cross_matrix.get_cross_matrix"], "The function uses the solution"

if __name__ == "__main__":
    os.environ["_PYTEST_RAISE"] = "1"
    pytest.main()


## Required packages

import os
import sys
import pytest
import numpy as np
import pandas as pd

import matplotlib as plt
plt.use('Agg')

sys.path.insert(0, os.path.abspath("../p_toolkit"))

from core import *

# -----------------------------------------------------------------------------
# p_adjust tests
# -----------------------------------------------------------------------------

def test_p_adjust_wrong_method_string():
    """
    Testing for entering an invalid string as a method.
    """

    try:
        d = [0.02,0.3]
        df = pd.DataFrame(data=d)
        p_adjust(data=df,  pv_index=0, method='bonferrro', alpha=0.05)
    except(ValueError):
        assert True
    else:
        assert False

def test_p_adjust_vector_1_value_bonf():
    """
    Testing p_adjust with a vector of 1 value and using bonferroni method.
    """

    d = {"p_value": [0.07], "adjusted": [0.07]}
    df = pd.DataFrame(data=d)
    df = df[["p_value", "adjusted"]]
    assert df.equals(p_adjust(data=[0.07], method="bonf")), "p_adjust 1 values vector for bonferoni"

def test_p_adjust_vector_1_value_bh():
    """
    Testing p_adjust with a vector of 1 value and using bh method.
    """

    d = {"p_value": [0.07], "adjusted": [0.07]}
    df = pd.DataFrame(data=d)
    df = df[["p_value", "adjusted"]]
    assert df.equals(p_adjust(data=[0.07], method="bh")), "p_adjust 1 values vector for bh"

def test_p_adjust_vector_2_values_bonf():
    """
    Testing p_adjust with a vector of 2 values and using bonferroni method.
    """

    d = {"p_value": [0.07, 0.2], "adjusted": [0.14, 0.4]}
    df = pd.DataFrame(data=d)
    df = df[["p_value", "adjusted"]]
    assert df.equals(p_adjust(data=[0.07, 0.2], method="bonf")), "p_adjust 2 values vector for bonferoni"

def test_p_adjust_vector_2_values_bh():
    """
    Testing p_adjust with a vector of 2 values and using bh method.
    """

    d = {"p_value": [0.07, 0.2], "adjusted": [0.14, 0.2]}
    df = pd.DataFrame(data=d)
    df = df[["p_value", "adjusted"]]
    assert df.equals(p_adjust(data=[0.07, 0.2], method="bh")), "p_adjust 2 values vector value for bh"

def test_p_adjust_character_col_index():
    """
    Testing for error string when col index of dataframe contains character values
    """

    try:
        err_str = {"p_value": ['str']}
        p_adjust((err_str), 0, "bonf", 0.05)
    except(TypeError):
        assert True
    else:
        assert False

def test_p_adjust_errors_probabilities_greater_than_one():
    """
    Testing for errors with invalid probabilities greater than one
    """

    try:
        err_str = {"p_value": [0.5,3,.02]}
        p_adjust((err_str), 0, "bonf", 0.05)
    except(TypeError):
        assert True
    else:
        assert False

def test_p_adjust_errors_probabilities_less_than_zero():
    """
    Testing for errors with invalid negative probabilities
    """

    try:
        err_str = {"p_value": [0.5,.3,-.02]}
        p_adjust((err_str), 0, "bh", 0.05)
    except(TypeError):
        assert True
    else:
        assert False

# -----------------------------------------------------------------------------
# p_methods tests
# -----------------------------------------------------------------------------

def test_p_methods_errors_probabilities_greater_than_one():
    """
    Testing for errors with invalid probabilities greater than one.
    """

    try:
        err_str = {"p_value": [0.5,3,.02]}
        p_methods((err_str), 0, 0.01)
    except(TypeError):
        assert True
    else:
        assert False

def test_p_methods_errors_probabilities_less_than_zero():
    """
    Testing for errors with invalid probabilities less than zero.
    """

    try:
        err_str = {"p_value": [0.5,.3,-.02]}
        p_methods((err_str), 0, 0.01)
    except(TypeError):
        assert True
    else:
        assert False

def test_p_methods_errors_alpha_less_than_zero():
    """
    Testing for errors with invalid alpha less than zero.
    """

    try:
        err_str = {"p_value": [0.5,.3,.02]}
        p_methods((err_str), 0, -.01)
    except(ProbabilityError):
        assert True
    else:
        assert False

def test_p_methods_errors_alpha_greater_than_one():
    """
    Testing for errors with invalid alpha greater than one.
    """

    try:
        err_str = {"p_value": [0.5,.3,.02]}
        p_methods((err_str), 0, 3)
    except(ProbabilityError):
        assert True
    else:
        assert False

def test_p_methods_1_value_vector_false_signficance():
    """
    Testing 1 value vector as input for p_methods with false significance.
    """

    d = {"p_value": [0.07], "bonf_value": [0.05], "bonf_significant": [False], "bh_value": [0.05],
         "bh_significant": [False]}
    df = pd.DataFrame(data=d)
    df = df[['p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    test = p_methods(data=[0.07], alpha=0.05)
    test = test[['p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    assert test.equals(df), "p_methods 1 value vector, FALSE"

def test_p_methods_1_value_vector_true_signficance():
    """
    Testing 1 value vector as input for p_methods with true significance.
    """

    d = {"p_value": [0.01], "bonf_value": [0.05], "bonf_significant": [True], "bh_value": [0.05],
         "bh_significant": [True]}
    df = pd.DataFrame(data=d)
    df = df[['p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    test = p_methods(data=[0.01], alpha=0.05)
    test = test[['p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    assert test.equals(df), "p_methods 1 value vector, TRUE"

def test_p_methods_2_value_vector():
    """
    Testing 2 values vector as input for p_methods.
    """

    d = {"Test": ["test 1", "test 2"], "p_value": [0.01, 0.03], "bonf_value": [0.025, 0.025],
         "bonf_significant": [True, False], "bh_value": [0.025, 0.05], "bh_significant": [True, True]}
    df = pd.DataFrame(data=d)
    df = df[['p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    test = p_methods(data=[0.01, 0.03], alpha=0.05)
    test = test[['p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    assert test.equals(df), "p_methods 2 values vector "

def test_p_methods_1_value_dataframe_true_signficance():
    """
    Testing 1 values dataframe with true signficance.
    """

    d = {"Test": ["test 1"], "p": [0.01]}
    df = pd.DataFrame(data=d)
    ad = {"Test": ["test 1"], "p_value": [0.01], "bonf_value": [0.05], "bonf_significant": [True], "bh_value": [0.05],
          "bh_significant": [True]}
    adf = pd.DataFrame(data=ad)
    adf = adf[['Test', 'p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    test = p_methods(data=df, pv_index="p", alpha=0.05)
    test = test[['Test', 'p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    assert test.equals(adf), "p_methods 2 values dataframe "

def test_p_methods_1_value_dataframe_false_signficance():
    """
    Testing 1 values dataframe with false signficance.
    """

    d = {"Test": ["test 1"], "p_value": [0.1]}
    df = pd.DataFrame(data=d)
    ad = {"Test": ["test 1"], "p_value": [0.1], "bonf_value": [0.05], "bonf_significant": [False], "bh_value": [0.05],
          "bh_significant": [False]}
    adf = pd.DataFrame(data=ad)
    adf = adf[['Test', 'p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    test = p_methods(data=df, pv_index="p_value", alpha=0.05)
    test = test[['Test', 'p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    assert test.equals(adf), "p_methods 2 values dataframe "

def test_p_methods_2_values_dataframe():
    """
    Testing 2 values dataframe.
    """

    d = {"Test": ["test 1", "test 2"], "p_value": [0.01, 0.03]}
    df = pd.DataFrame(data=d)
    ad = {"Test": ["test 1", "test 2"], "p_value": [0.01, 0.03], "bonf_value": [0.025, 0.025],
          "bonf_significant": [True, False], "bh_value": [0.025, 0.05], "bh_significant": [True, True]}
    adf = pd.DataFrame(data=ad)
    adf = adf[['Test', 'p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    test = p_methods(data=df, pv_index="p_value", alpha=0.05)
    test = test[['Test', 'p_value', 'bh_value', 'bh_significant', 'bonf_value', 'bonf_significant']]
    assert test.equals(adf), "p_methods 2 values dataframe "

def test_p_methods_empty_data():
    """
    Testing not empty data.
    """
    with pytest.raises(TypeError):
        p_methods()

# -----------------------------------------------------------------------------
# p_qq tests
# -----------------------------------------------------------------------------

def test_p_qq_error_p_value_strings():
    """
    Testing for errors with p_qq with strings in the p_values vector.
    """

    try:
        err_str = {"p_value": ['str']}
        p_qq((err_str), 0, "bonf", 0.05)
    except(TypeError):
        assert True
    else:
        assert False

def test_p_qq_errors_probabilities_less_than_zero():
    """
    Testing for errors with probabilities less than zero.
    """
    try:
        err_str = {"p_value": [0.5,3,.02]}
        p_qq((err_str), 0, 0.01)
    except(TypeError):
        assert True
    else:
        assert False

def test_p_qq_errors_probabilities_greater_than_one():
    """
    Testing for errors with probatilities greater than one.
    """

    try:
        err_str = {"p_value": [0.5,.3,-.02]}
        p_qq((err_str), 0, 0.01)
    except(TypeError):
        assert True
    else:
        assert False

# -----------------------------------------------------------------------------
# p_plot tests
# -----------------------------------------------------------------------------

def test_p_plot_errors_probabilities_greater_than_one():
    """
    Testing for errors with invalid probabilities greater than one.
    """
    try:
        err_str = {"p_value": [0.5,3,.02]}
        p_plot((err_str), 0, 0.01)
    except(TypeError):
        assert True
    else:
        assert False

def test_p_plot_errors_probabilities_less_than_zero():
    """
    Testing for errors with invalid negative probabilities.
    """

    try:
        err_str = {"p_value": [0.5,.3,-.02]}
        p_plot((err_str), 0, 0.01)
    except(TypeError):
        assert True
    else:
        assert False

def test_p_plot_errors_alpha_less_than_zero():
    """
    Testing for errors with invalid negative alpha.
    """
    try:
        err_str = {"p_value": [0.5,.3,.02]}
        p_plot((err_str), 0, -.01)
    except(TypeError):
        assert True
    else:
        assert False

def test_p_plot_errors_alpha_greater_than_one():
    """
    Testing for errors with invalid alpha greater than one.
    """
    try:
        err_str = {"p_value": [0.5,.3,.02]}
        p_plot((err_str), 0, 3)
    except(TypeError):
        assert True
    else:
        assert False

# -----------------------------------------------------------------------------
# Integration tests
# -----------------------------------------------------------------------------

def test_p_plot_integration_test():
    """
    Integration test using the p_plot function with p_methods as input.
    """
    try:
        p_plot(p_methods(data=[0.01], alpha=0.05))
    except(SyntaxError):
        assert False
    else:
        assert True

def test_p_qq_integration_test():
    """
    Integration test using the p_qq function with p_methods as input.
    """
    try:
        p_qq(p_methods(data=[0.01], alpha=0.05))
    except(SyntaxError):
        assert False
    else:
        assert True


#!usr/bin/env python3
import numba
from numba.typed import List

# JIT, Just-In-Time compiling, is a technology that has been used very
# successfully in several dynamic languages to gain performance.
# The idea behind Just-In-Time is that we collect information on the code at
# run time and according to this information, we generate specific machine
# code for these cases. If you're function runs only once, JIT will not help.
# But in most cases, bottlenecks are in functions that are called many times.
# Numba, this is a compiler kit called LLVM, under the hood, to generate
# machine code for Python functions.
# LLVM is an acronym that stands for low level virtual machine.
# It also refers to a compiling technology called the LLVM project, which is a
# collection of modular and reusable compiler and toolchain technologies.

def poly(coeffs, n):
    """Compute value of polynomial given coefficients"""
    total = 0
    for i, c in enumerate(coeffs):
        total += c * n ** i
    return total

@numba.jit
def poly_j(coeffs, n):
    """Compute value of polynomial given coefficients - JIT"""
    total = 0
    for i, c in enumerate(coeffs):
        total += c * n ** i
    return total

if __name__ == '__main__':
    # previously we had coeffs = [4, 8, 15, 16, 23, 42] but we were getting
    # a warning:
    # http://numba.pydata.org/numba-doc/latest/reference/deprecation.html#deprecation-of-reflection-for-list-and-set-types
    coeffs = List()
    [coeffs.append(x) for x in [4, 8, 15, 16, 23, 42]]

# In [7]: %run src/using_jit.py
#
# In [8]: poly(coeffs, 7)
# Out[8]: 767400
#
# In [9]: poly_j(coeffs, 7)
# Out[9]: 767400
#
# In [10]: %timeit poly(coeffs, 7)
# 18.6 µs ± 2.08 µs per loop (mean ± std. dev. of 7 runs, 100000 loops each)
#
# In [11]: %timeit poly_j(coeffs, 7)
# 1.67 µs ± 36.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)


"""
.. _tut_io_export_pandas:

=================================
Export epochs to Pandas DataFrame
=================================

In this example the pandas exporter will be used to produce a DataFrame
object. After exploring some basic features a split-apply-combine
work flow will be conducted to examine the latencies of the response
maxima across epochs and conditions.
Note. Equivalent methods are available for raw and evoked data objects.

Short Pandas Primer
-------------------

Pandas Data Frames
~~~~~~~~~~~~~~~~~~
A data frame can be thought of as a combination of matrix, list and dict:
It knows about linear algebra and element-wise operations but is size mutable
and allows for labeled access to its data. In addition, the pandas data frame
class provides many useful methods for restructuring, reshaping and visualizing
data. As most methods return data frame instances, operations can be chained
with ease; this allows to write efficient one-liners. Technically a DataFrame
can be seen as a high-level container for numpy arrays and hence switching
back and forth between numpy arrays and DataFrames is very easy.
Taken together, these features qualify data frames for inter operation with
databases and for interactive data exploration / analysis.
Additionally, pandas interfaces with the R statistical computing language that
covers a huge amount of statistical functionality.

Export Options
~~~~~~~~~~~~~~
The pandas exporter comes with a few options worth being commented.

Pandas DataFrame objects use a so called hierarchical index. This can be
thought of as an array of unique tuples, in our case, representing the higher
dimensional MEG data in a 2D data table. The column names are the channel names
from the epoch object. The channels can be accessed like entries of a
dictionary:

    df['MEG 2333']

Epochs and time slices can be accessed with the .ix method:

    epochs_df.ix[(1, 2), 'MEG 2333']

However, it is also possible to include this index as regular categorial data
columns which yields a long table format typically used for repeated measure
designs. To take control of this feature, on export, you can specify which
of the three dimensions 'condition', 'epoch' and 'time' is passed to the Pandas
index using the index parameter. Note that this decision is revertible any
time, as demonstrated below.

Similarly, for convenience, it is possible to scale the times, e.g. from
seconds to milliseconds.

Some Instance Methods
~~~~~~~~~~~~~~~~~~~~~
Most numpy methods and many ufuncs can be found as instance methods, e.g.
mean, median, var, std, mul, , max, argmax etc.
Below an incomplete listing of additional useful data frame instance methods:

apply : apply function to data.
    Any kind of custom function can be applied to the data. In combination with
    lambda this can be very useful.
describe : quickly generate summary stats
    Very useful for exploring data.
groupby : generate subgroups and initialize a 'split-apply-combine' operation.
    Creates a group object. Subsequently, methods like apply, agg, or transform
    can be used to manipulate the underlying data separately but
    simultaneously. Finally, reset_index can be used to combine the results
    back into a data frame.
plot : wrapper around plt.plot
    However it comes with some special options. For examples see below.
shape : shape attribute
    gets the dimensions of the data frame.
values :
    return underlying numpy array.
to_records :
    export data as numpy record array.
to_dict :
    export data as dict of arrays.

Reference
~~~~~~~~~
More information and additional introductory materials can be found at the
pandas doc sites: http://pandas.pydata.org/pandas-docs/stable/
"""
# Author: Denis Engemann <denis.engemann@gmail.com>
#
# License: BSD (3-clause)

import mne
import matplotlib.pyplot as plt
import numpy as np
from mne.datasets import sample

print(__doc__)

data_path = sample.data_path()
raw_fname = data_path + '/MEG/sample/sample_audvis_filt-0-40_raw.fif'
event_fname = data_path + '/MEG/sample/sample_audvis_filt-0-40_raw-eve.fif'

raw = mne.io.read_raw_fif(raw_fname)

# For simplicity we will only consider the first 10 epochs
events = mne.read_events(event_fname)[:10]

# Add a bad channel
raw.info['bads'] += ['MEG 2443']
picks = mne.pick_types(raw.info, meg='grad', eeg=False, eog=True,
                       stim=False, exclude='bads')

tmin, tmax = -0.2, 0.5
baseline = (None, 0)
reject = dict(grad=4000e-13, eog=150e-6)

event_id = dict(auditory_l=1, auditory_r=2, visual_l=3, visual_r=4)

epochs = mne.Epochs(raw, events, event_id, tmin, tmax, proj=True, picks=picks,
                    baseline=baseline, preload=True, reject=reject)

###############################################################################
# Export DataFrame

# The following parameters will scale the channels and times plotting
# friendly. The info columns 'epoch' and 'time' will be used as hierarchical
# index whereas the condition is treated as categorial data. Note that
# this is optional. By passing None you could also print out all nesting
# factors in a long table style commonly used for analyzing repeated measure
# designs.

index, scale_time, scalings = ['epoch', 'time'], 1e3, dict(grad=1e13)

df = epochs.to_data_frame(picks=None, scalings=scalings, scale_time=scale_time,
                          index=index)

# Create MEG channel selector and drop EOG channel.
meg_chs = [c for c in df.columns if 'MEG' in c]

df.pop('EOG 061')  # this works just like with a list.

###############################################################################
# Explore Pandas MultiIndex

# Pandas is using a MultiIndex or hierarchical index to handle higher
# dimensionality while at the same time representing data in a flat 2d manner.

print(df.index.names, df.index.levels)

# Inspecting the index object unveils that 'epoch', 'time' are used
# for subsetting data. We can take advantage of that by using the
# .ix attribute, where in this case the first position indexes the MultiIndex
# and the second the columns, that is, channels.

# Plot some channels across the first three epochs
xticks, sel = np.arange(3, 600, 120), meg_chs[:15]
df.ix[:3, sel].plot(xticks=xticks)
mne.viz.tight_layout()

# slice the time starting at t0 in epoch 2 and ending 500ms after
# the base line in epoch 3. Note that the second part of the tuple
# represents time in milliseconds from stimulus onset.
df.ix[(1, 0):(3, 500), sel].plot(xticks=xticks)
mne.viz.tight_layout()

# Note: For convenience the index was converted from floating point values
# to integer values. To restore the original values you can e.g. say
# df['times'] = np.tile(epoch.times, len(epochs_times)

# We now reset the index of the DataFrame to expose some Pandas
# pivoting functionality. To simplify the groupby operation we
# we drop the indices to treat epoch and time as categroial factors.

df = df.reset_index()

# The ensuing DataFrame then is split into subsets reflecting a crossing
# between condition and trial number. The idea is that we can broadcast
# operations into each cell simultaneously.

factors = ['condition', 'epoch']
sel = factors + ['MEG 1332', 'MEG 1342']
grouped = df[sel].groupby(factors)

# To make the plot labels more readable let's edit the values of 'condition'.
df.condition = df.condition.apply(lambda name: name + ' ')

# Now we compare the mean of two channels response across conditions.
grouped.mean().plot(kind='bar', stacked=True, title='Mean MEG Response',
                    color=['steelblue', 'orange'])
mne.viz.tight_layout()

# We can even accomplish more complicated tasks in a few lines calling
# apply method and passing a function. Assume we wanted to know the time
# slice of the maximum response for each condition.

max_latency = grouped[sel[2]].apply(lambda x: df.time[x.argmax()])

print(max_latency)

# Then make the plot labels more readable let's edit the values of 'condition'.
df.condition = df.condition.apply(lambda name: name + ' ')

plt.figure()
max_latency.plot(kind='barh', title='Latency of Maximum Response',
                 color=['steelblue'])
mne.viz.tight_layout()

# Finally, we will again remove the index to create a proper data table that
# can be used with statistical packages like statsmodels or R.

final_df = max_latency.reset_index()
final_df.rename(columns={0: sel[2]})  # as the index is oblivious of names.

# The index is now written into regular columns so it can be used as factor.
print(final_df)

plt.show()

# To save as csv file, uncomment the next line.
# final_df.to_csv('my_epochs.csv')

# Note. Data Frames can be easily concatenated, e.g., across subjects.
# E.g. say:
#
# import pandas as pd
# group = pd.concat([df_1, df_2])
# group['subject'] = np.r_[np.ones(len(df_1)), np.ones(len(df_2)) + 1]


""" Helper Functions Assignment """

import pandas as pd
import numpy as np

def null_count(df):
    # return the amount of null values in a dataframe
    return df.isnull().sum()

def randomized(df, seed):
    np.random.seed(seed)
    df = df.iloc[np.random.shuffle(len(df))].reset_index(drop=True)
    return df


# Julian Conneely 05/10/18
# adding Title, Labels and Legend

import matplotlib.pyplot as plt
import numpy as np

x = np.arange(0.0, 10.0, 0.01)
y = 3.0 * x + 1.0
noise = np.random.normal(0.0, 1.0, len(x))

plt.plot(x, y + noise, 'r.', label="Actual") # added label
plt.plot(x, y, 'b-', label="Model") # added label

plt.title("Super Simple Plot") # title, labels and legend
plt.xlabel("weight")
plt.ylabel("height")
plt.legend()

plt.show()


"""
Generating a CCD mosaic
=======================

This file contains a class to create a single VIS CCD image from separate files one for each quadrant.

:requires: NumPy
:requires: PyFITS

:author: Sami-Matias Niemi
:contact: s.niemi@ucl.ac.uk

To execute::

    python tileCCD.py -f 'Q*science.fits' -e 1

where -f argument defines the input files to be tiled and the -e argument marks the
FITS extension from which the imaging data are being read.

:version: 0.5

.. todo::

    #. Does not deal properly with multiple WCSs coming in the different quadrants (should recalculate the
       centre of the CCD and modify the WCS accordingly).
    #. Improve the history section.

"""
import pyfits as pf
import numpy as np
from optparse import OptionParser
import sys, os, datetime, re
import glob as g
from support import logger as lg

class tileCCD():
    """
    Class to create a single VIS CCD image from separate quadrants files.
    """
    def __init__(self, inputs, log):
        """
        Class constructor.
        """
        self.inputs = inputs
        self.log = log

    def readData(self):
        """
        Reads in data from all the input files and the header from the first file.
        Input files are taken from the input dictionary given when class was initiated.

        Subtracts the pre- and overscan regions if these were simulated. Takes into account
        which quadrant is being processed so that the extra regions are subtracted correctly.
        """
        data = {}
        for i, file in enumerate(self.inputs['files']):
            fh = pf.open(file)
            hdu = fh[self.inputs['ext']].header

            if i == 0:
                self.hdu = hdu

            try:
                overscan = hdu['OVERSCA']
            except:
                overscan = 'False'

            if 'True' in overscan and 'nonoise' not in file:
                self.log.info('Subtracting pre- and overscan regions')
                prescanx = hdu['PRESCANX']
                overscanx = hdu['OVRSCANX']
                quadrant = hdu['QUADRANT']

                if quadrant in (0, 2):
                    data[file] = fh[self.inputs['ext']].data[:, prescanx: -overscanx]
                else:
                    data[file] = fh[self.inputs['ext']].data[:, overscanx: -prescanx]
            else:
                self.log.info('No overscan simulated, using full array')
                data[file] = fh[self.inputs['ext']].data

            self.log.info('Read data from {0:>s} extension {1:d}'.format(file, self.inputs['ext']))
            fh.close()

        self.data = data
        return self.data

    def tileCCD(self, xsize=2048, ysize=2066):
        """
        Tiles quadrants to form a single CCD image.

        Assume that the input file naming convention is Qx_CCDX_CCDY_name.fits.

        :param xsize: length of a quadrant in column direction
        :type xsize: int
        :param ysize: length of a quadrant in row direction
        :type ysize: int

        :return: image array of size (ysize*2, xsize*2)
        :rtype: ndnarray
        """
        self.CCDdata = np.zeros((ysize*2, xsize*2))

        for key, data in self.data.iteritems():
            #use regular expression to find the numbers
            p = re.compile('\d+')
            ls = p.findall(key)
            #in ls, numbers are [quadrant, CCDx, CCDy]
            if len(ls) < 3:
                print 'Problem when parsing the file name!'
                print 'Filenames should be in format:'
                print 'Qx_CCDX_CCDYfilename.fits'
                self.log.error('Problem when parsing the file name!')
                return self.CCDdata

            #note here that because we will save a FITS file the numbering
            #starts from the lower left corner!
            if int(ls[0]) == 0:
                self.CCDdata[0:ysize, 0:xsize] = data
            elif int(ls[0]) == 1:
                self.CCDdata[:ysize, xsize:] = data
            elif int(ls[0]) == 2:
                self.CCDdata[ysize:, :xsize] = data
            elif int(ls[0]) == 3:
                self.CCDdata[ysize:, xsize:] = data

        return self.CCDdata

    def writeFITSfile(self, data=None, unsigned16bit=True):
        """
        Write out FITS files using PyFITS.

        :param data: data to write to a FITS file, if None use self.data
        :type data: ndarray
        :param unsigned16bit: whether to scale the data using bzero=32768
        :type unsigned16bit: bool

        :return: None
        """
        if os.path.isfile(self.inputs['output']):
            self.log.info('Deleted existing file %s to generate a new output' % self.inputs['output'])
            os.remove(self.inputs['output'])
        else:
            self.log.info('Writing output to %s' % self.inputs['output'])

        if data is None:
            data = self.CCDdata

        #create a new FITS file, using HDUList instance
        ofd = pf.HDUList()

        #new image HDU
        hdu = pf.PrimaryHDU(data, self.hdu)

        #convert to unsigned 16bit int if requested
        if unsigned16bit:
            hdu.scale('int16', '', bzero=32768)
            hdu.header.add_history('Scaled to unsigned 16bit integer!')

        #add keywords
        for key, value in self.inputs.iteritems():
            try:
                hdu.header.add_history('{0:>s} = {1:>s}'.format(key, value))
            except:
                hdu.header.add_history('{0:>s} = {1:d}'.format(key, value))

        #update and verify the header
        hdu.header.add_history('If questions, please contact Sami-Matias Niemi (s.niemi at ucl.ac.uk).')
        hdu.header.add_history('This file has been created with the VISsim Python Package at %s' \
                               % datetime.datetime.isoformat(datetime.datetime.now()))
        hdu.verify('fix')

        ofd.append(hdu)

        #write the actual file
        ofd.writeto(self.inputs['output'])
        self.log.info('Wrote %s' % self.inputs['output'])

    def runAll(self):
        """
        Wrapper to perform all class methods.
        """
        self.readData()
        self.tileCCD()
        self.writeFITSfile()

def processArgs(printHelp=False):
    """
    Processes command line arguments.
    """
    parser = OptionParser()

    parser.add_option('-f', '--files', dest='files',
                      help="Input files to compile e.g. 'Q*_00_00science.fits'", metavar='string')
    parser.add_option('-o', '--output', dest='output',
                      help="Name of the output file, default=VISCCD.fits", metavar='string')
    parser.add_option('-e', '--extension', type='int', dest='ext',
                     help='FITS extension from which to look for data, default=0', metavar='int')
    parser.add_option('-d', '--debug', dest='debug', action='store_true',
                      help='Debugging mode on')
    if printHelp:
        parser.print_help()
    else:
        return parser.parse_args()

if __name__ == '__main__':
    opts, args = processArgs()

    if opts.files is None:
        processArgs(True)
        sys.exit(8)

    #FITS extension
    if opts.ext is None:
        ext = 0
    else:
        ext = opts.ext

    #name of the output file
    if opts.output is None:
        output = 'VISCCD.fits'
    else:
        output = opts.output

    #logger
    log = lg.setUpLogger('tileCCDs.log')

    #look for files
    files = g.glob(opts.files)
    files.sort()
    if len(files) / 4. > 1.0 or len(files) == 0:
        print 'Detected %i input files, but the current version does not support anything but tiling four files...' \
              % len(files)
        sys.exit(9)

    #write to the log what files were used
    log.info('Input files:')
    for file in files:
        log.info(file)

    #intputs
    inputs = dict(files=files, ext=ext, output=output)

    #class call
    tile = tileCCD(inputs, log)
    tile.runAll()

    log.info('CCD tiled, script will exit...')