import json
import os
import logging
import tempfile

import numpy as np
import meds.util
import fitsio
import galsim
import yaml
from esutil.ostools import StagedOutFile
from meds.maker import MEDSMaker
import desmeds.util

from .constants import MAGZP_REF, MEDSCONF
from .psf_wrapper import PSFWrapper
from .wcsing import get_galsim_wcs
from .files import (
    get_band_info_file, get_meds_file_path, expand_path, make_dirs_for_file)

logger = logging.getLogger(__name__)

def make_meds_files(*, tilename, bands, output_meds_dir, psf_kws, meds_config):
    """Make a MEDS file for a given band and tilename.

    Parameters
    ----------
    tilename : str
        The DES coadd tile to run true detection.
    bands : str
        The bands to run true detection.
    output_meds_dir : str
        The output DEADATA/MEDS_DIR for the simulation data products.
    psf_kws : dict
        The dictionary of PSF config information.
    meds_config : dict
        The MEDS making configuration file. See the default one in
        `work/simple_des_y3_sims/default_configs/meds.yaml`.
    """

    logger.info(' making meds files for coadd tile %s', tilename)

    # force this
    meds_config['magzp_ref'] = MAGZP_REF
    meds_config['psf'] = {'type': 'psfex'}
    meds_config['use_joblib'] = True

    # read info files
    info = {}
    for band in bands:
        # get info about files
        fname = get_band_info_file(
            meds_dir=output_meds_dir,
            medsconf=MEDSCONF,
            tilename=tilename,
            band=band)
        with open(fname, 'r') as fp:
            info[band] = yaml.load(fp, Loader=yaml.Loader)

    # always get the truth catalog from r band
    cat = fitsio.read(info['r']['cat_path'].replace(
        '$MEDS_DIR', output_meds_dir))

    for band in bands:
        logger.info(' doing band %s', band)

        # get all of the components for the file
        obj_data = _make_meds_input_data_struct(
            cat=cat,
            allowed_box_sizes=meds_config['allowed_box_sizes'],
            min_box_size=meds_config['min_box_size'],
            max_box_size=meds_config['max_box_size'],
            sigma_fac=meds_config['sigma_fac'])
        image_info = _make_meds_image_info_struct(
            info=info[band], output_meds_dir=output_meds_dir)
        meta_data = _make_meds_metadata(band=band, tilename=tilename)
        psf_data = _build_psf_data(
            info=info[band],
            psf_kws=psf_kws,
            output_meds_dir=output_meds_dir)

        # make the file in a tmp dir and then stage out
        maker = MEDSMaker(
            obj_data,
            image_info,
            psf_data=psf_data,
            config=meds_config,
            meta_data=meta_data)

        final_meds_file = get_meds_file_path(
            meds_dir=output_meds_dir,
            medsconf=MEDSCONF,
            tilename=tilename,
            band=band)
        make_dirs_for_file(final_meds_file)

        with tempfile.TemporaryDirectory() as tmpdir:
            with StagedOutFile(final_meds_file, tmpdir=tmpdir) as sf:
                uncompressed_file = sf.path.replace('.fits.fz', '.fits')
                make_dirs_for_file(uncompressed_file)
                maker.write(uncompressed_file)

                # make sure to remove the destination file when fpacking
                try:
                    os.remove(sf.path)
                except Exception:
                    pass
                desmeds.util.fpack_file(uncompressed_file)
                try:
                    os.remove(uncompressed_file)
                except Exception:
                    pass

def _build_psf_data(*, info, psf_kws, output_meds_dir):
    def _load_psf_data(_info, force_gauss=False):
        wcs = get_galsim_wcs(
            image_path=_info['image_path'].replace(
                '$MEDS_DIR', output_meds_dir),
            image_ext=_info['image_ext'])
        if psf_kws['type'] == 'gauss' or force_gauss:
            return PSFWrapper(galsim.Gaussian(fwhm=0.9), wcs)
        elif psf_kws['type'] == 'piff':
            from ..des_piff import DES_Piff
            piff_model = DES_Piff(expand_path(_info['piff_path']))
            return PSFWrapper(piff_model, wcs)
        elif psf_kws['type'] == 'gauss-pix':
            from .gauss_pix_psf import GaussPixPSF
            kwargs = {k: psf_kws[k] for k in psf_kws if k != 'type'}
            psf_model = GaussPixPSF(**kwargs)
            return PSFWrapper(psf_model, wcs)
        else:
            raise ValueError("psf type '%s' is not valid!" % psf_kws['type'])

    force_gauss = psf_kws['type'] in ['piff']
    psf_data = [_load_psf_data(info, force_gauss=force_gauss)]
    for se_info in info['src_info']:
        psf_data.append(_load_psf_data(se_info))
    return psf_data

def _make_meds_metadata(*, band, tilename):
    meta = np.zeros(1, dtype=[
        ('magzp_ref', 'f8'),
        ('band', 'S1'),
        ('tilename', 'S12')])
    meta['magzp_ref'] = MAGZP_REF
    meta['band'] = band
    meta['tilename'] = tilename
    return meta

def _make_meds_image_info_struct(*, info, output_meds_dir):
    def _munge_path(pth):
        return pth.replace('$MEDS_DIR', output_meds_dir)

    # get WCS structures
    wcs_json = _load_wcs_json(info=info, output_meds_dir=output_meds_dir)
    wcs_len = max([len(j) for j in wcs_json])

    # compute the max path length
    path_len = [
        len(_munge_path(info['image_path'])),
        len(_munge_path(info['weight_path'])),
        len(_munge_path(info['seg_path'])),
        len(_munge_path(info['bmask_path']))]
    for se_info in info['src_info']:
        path_len += [
            len(_munge_path(se_info['image_path'])),
            len(_munge_path(se_info['weight_path'])),
            len(_munge_path(se_info['bkg_path'])),
            len(_munge_path(se_info['bmask_path']))]
    path_len = max(path_len)

    # now fill the array
    dtype = meds.util.get_image_info_dtype(
        path_len,
        wcs_len=wcs_len,
        ext_len=3)
    image_info = np.zeros(len(info['src_info']) + 1, dtype=dtype)

    image_info['image_id'] = np.arange(len(image_info))
    image_info['image_flags'] = 0
    image_info['position_offset'] = 1
    for i, wj in enumerate(wcs_json):
        image_info['wcs'][i] = wj

    image_info['scale'][0] = info['scale']
    image_info['magzp'][0] = info['magzp']
    image_info['image_path'][0] = _munge_path(info['image_path'])
    image_info['image_ext'][0] = info['image_ext']
    image_info['weight_path'][0] = _munge_path(info['weight_path'])
    image_info['weight_ext'][0] = info['weight_ext']
    image_info['bmask_path'][0] = _munge_path(info['bmask_path'])
    image_info['bmask_ext'][0] = info['bmask_ext']
    image_info['bkg_path'][0] = ""
    image_info['bkg_ext'][0] = ""
    image_info['seg_path'][0] = _munge_path(info['seg_path'])
    image_info['seg_ext'][0] = info['seg_ext']
    for i, se_info in enumerate(info['src_info']):
        image_info['scale'][i+1] = se_info['scale']
        image_info['magzp'][i+1] = se_info['magzp']
        image_info['image_path'][i+1] = _munge_path(se_info['image_path'])
        image_info['image_ext'][i+1] = se_info['image_ext']
        image_info['weight_path'][i+1] = _munge_path(se_info['weight_path'])
        image_info['weight_ext'][i+1] = se_info['weight_ext']
        image_info['bmask_path'][i+1] = _munge_path(se_info['bmask_path'])
        image_info['bmask_ext'][i+1] = se_info['bmask_ext']
        image_info['bkg_path'][i+1] = _munge_path(se_info['bkg_path'])
        image_info['bkg_ext'][i+1] = se_info['bkg_ext']
        image_info['seg_path'][i+1] = ""
        image_info['seg_ext'][i+1] = ""

    return image_info

def _load_wcs_json(*, info, output_meds_dir):
    def _munge_header(hd):
        return {k.lower(): hd[k] for k in hd if k is not None}

    wcs_json = []
    hd = fitsio.read_header(
        info['image_path'].replace('$MEDS_DIR', output_meds_dir),
        ext=info['image_ext'])
    wcs_json.append(json.dumps(_munge_header(hd)))
    for se_info in info['src_info']:
        hd = fitsio.read_header(
            se_info['image_path'].replace('$MEDS_DIR', output_meds_dir),
            ext=se_info['image_ext'])
        wcs_json.append(json.dumps(_munge_header(hd)))

    return wcs_json

def _make_meds_input_data_struct(
        *, cat, allowed_box_sizes, min_box_size, max_box_size, sigma_fac):
    """Make the input data structure for the MEDS maker.

    Parameters
    ----------
    cat : np.ndarray
        The coadd catalog.
    allowed_box_sizes : list of ints
        A list of the allowed postage stamp box sizes.
    min_box_size : int
        The minimum allowed box size. This value overrides any smaller values
        in `allowed_box_sizes`.
    max_box_size : int
        The maximum allowed box size. This value overrides any larger values
        in `allowed_box_sizes`.
    sigma_fac : float
        The factor by which to scale the flux radius. A value around 5 is
        standard.

    Returns
    -------
    input_data : np.ndarray
        The input structured array for the MEDS maker.
    """
    dtype = meds.util.get_meds_input_dtype(extra_fields=[('number', 'i8')])
    input_data = np.zeros(len(cat), dtype=dtype)
    input_data['id'] = cat['number']
    input_data['number'] = cat['number']
    input_data['ra'] = cat['alpha_j2000']
    input_data['dec'] = cat['delta_j2000']
    input_data['box_size'] = _get_box_sizes(
        cat=cat,
        allowed_box_sizes=allowed_box_sizes,
        min_box_size=min_box_size,
        max_box_size=max_box_size,
        sigma_fac=sigma_fac)
    return input_data

def _get_box_sizes(
        *, cat, allowed_box_sizes, min_box_size, max_box_size, sigma_fac):
    """Get the box sizes for the coadd catalog.

    Parameters
    ----------
    cat : np.ndarray
        The coadd catalog.
    allowed_box_sizes : list of ints
        A list of the allowed postage stamp box sizes.
    min_box_size : int
        The minimum allowed box size. This value overrides any smaller values
        in `allowed_box_sizes`.
    max_box_size : int
        The maximum allowed box size. This value overrides any larger values
        in `allowed_box_sizes`.
    sigma_fac : float
        The factor by which to scale the flux radius. A value around 5 is
        standard.

    Returns
    -------
    box_sizes : np.ndarray
        The array of box sizes.
    """
    sigma_size = get_sigma_size(cat=cat, sigma_fac=sigma_fac)

    # now do row and col sizes
    row_size = cat['ymax_image'] - cat['ymin_image'] + 1
    col_size = cat['xmax_image'] - cat['xmin_image'] + 1

    # get max of all three
    box_size = np.vstack(
        (col_size, row_size, sigma_size)).max(axis=0)

    # clip to range
    box_size = box_size.clip(min_box_size, max_box_size)

    # now put in fft sizes
    bins = [0]
    bins.extend([sze for sze in allowed_box_sizes
                 if sze >= min_box_size and sze <= max_box_size])

    if bins[-1] != max_box_size:
        bins.append(max_box_size)

    bin_inds = np.digitize(box_size, bins, right=True)
    bins = np.array(bins)

    return bins[bin_inds]

def get_sigma_size(*, cat, sigma_fac):
    """Get an object size based on its flux radius and ellipticity.

    Parameters
    ----------
    cat : np.ndarray
        The coadd catalog.
    sigma_fac : float
        The factor by which to scale the flux radius. A value around 5 is
        standard.

    Returns
    -------
    sigma_size : np.ndarray
        The array of sizes.
    """

    fwhm_fac = 2*np.sqrt(2*np.log(2))

    ellipticity = 1.0 - cat['b_world'] / cat['a_world']
    sigma = cat['flux_radius'] * 2.0 / fwhm_fac
    drad = sigma * sigma_fac
    drad = drad * (1.0 + ellipticity)
    drad = np.ceil(drad)
    # sigma size is twice the radius
    sigma_size = 2 * drad.astype('i4')

    return sigma_size


"""Supporting functions for the 'reference' command."""
import collections
import logging

import numpy as np
import pandas as pd
import pyfaidx
from skgenome import tabio, GenomicArray as GA

from . import core, fix, descriptives, params
from .cmdutil import read_cna
from .cnary import CopyNumArray as CNA

def do_reference_flat(targets, antitargets=None, fa_fname=None,
                      male_reference=False):
    """Compile a neutral-coverage reference from the given intervals.

    Combines the intervals, shifts chrX values if requested, and calculates GC
    and RepeatMasker content from the genome FASTA sequence.
    """
    ref_probes = bed2probes(targets)
    if antitargets:
        ref_probes.add(bed2probes(antitargets))
    # Set sex chromosomes by "reference" sex
    ref_probes['log2'] = ref_probes.expect_flat_log2(male_reference)
    ref_probes['depth'] = np.exp2(ref_probes['log2'])  # Shim
    # Calculate GC and RepeatMasker content for each probe's genomic region
    if fa_fname:
        gc, rmask = get_fasta_stats(ref_probes, fa_fname)
        ref_probes['gc'] = gc
        ref_probes['rmask'] = rmask
    else:
        logging.info("No FASTA reference genome provided; "
                     "skipping GC, RM calculations")
    ref_probes.sort_columns()
    return ref_probes

def bed2probes(bed_fname):
    """Create a neutral-coverage CopyNumArray from a file of regions."""
    regions = tabio.read_auto(bed_fname)
    table = regions.data.loc[:, ("chromosome", "start", "end")]
    table["gene"] = (regions.data["gene"] if "gene" in regions.data else '-')
    table["log2"] = 0.0
    table["spread"] = 0.0
    return CNA(table, {"sample_id": core.fbase(bed_fname)})

def do_reference(target_fnames, antitarget_fnames=None, fa_fname=None,
                 male_reference=False, female_samples=None,
                 do_gc=True, do_edge=True, do_rmask=True, do_cluster=False,
                 min_cluster_size=4):
    """Compile a coverage reference from the given files (normal samples)."""
    if antitarget_fnames:
        core.assert_equal("Unequal number of target and antitarget files given",
                          targets=len(target_fnames),
                          antitargets=len(antitarget_fnames))
    if not fa_fname:
        logging.info("No FASTA reference genome provided; "
                     "skipping GC, RM calculations")

    if female_samples is None:
        # NB: Antitargets are usually preferred for inferring sex, but might be
        # empty files, in which case no inference can be done. Since targets are
        # guaranteed to exist, infer from those first, then replace those
        # values where antitargets are suitable.
        sexes = infer_sexes(target_fnames, False)
        if antitarget_fnames:
            a_sexes = infer_sexes(antitarget_fnames, False)
            for sid, a_is_xx in a_sexes.items():
                t_is_xx = sexes.get(sid)
                if t_is_xx is None:
                    sexes[sid] = a_is_xx
                elif t_is_xx != a_is_xx and a_is_xx is not None:
                    logging.warning("Sample %s chromosomal X/Y ploidy looks "
                                    "like %s in targets but %s in antitargets; "
                                    "preferring antitargets",
                                    sid,
                                    "female" if t_is_xx else "male",
                                    "female" if a_is_xx else "male")
                    sexes[sid] = a_is_xx
    else:
        sexes = collections.defaultdict(lambda: female_samples)

    # TODO - refactor/inline this func here, once it works
    ref_probes = combine_probes(target_fnames, antitarget_fnames, fa_fname,
                                male_reference, sexes, do_gc, do_edge, do_rmask,
                                do_cluster, min_cluster_size)
    warn_bad_bins(ref_probes)
    return ref_probes

def infer_sexes(cnn_fnames, is_haploid_x):
    """Map sample IDs to inferred chromosomal sex, where possible.

    For samples where the source file is empty or does not include either sex
    chromosome, that sample ID will not be in the returned dictionary.
    """
    sexes = {}
    for fname in cnn_fnames:
        cnarr = read_cna(fname)
        if cnarr:
            is_xx = cnarr.guess_xx(is_haploid_x)
            if is_xx is not None:
                sexes[cnarr.sample_id] = is_xx
    return sexes

def combine_probes(filenames, antitarget_fnames, fa_fname,
                   is_haploid_x, sexes, fix_gc, fix_edge, fix_rmask,
                   do_cluster, min_cluster_size):
    """Calculate the median coverage of each bin across multiple samples.

    Parameters
    ----------
    filenames : list
        List of string filenames, corresponding to targetcoverage.cnn and
        antitargetcoverage.cnn files, as generated by 'coverage' or
        'import-picard'.
    fa_fname : str
        Reference genome sequence in FASTA format, used to extract GC and
        RepeatMasker content of each genomic bin.
    is_haploid_x : bool
    do_cluster : bool
    fix_gc : bool
    fix_edge : bool
    fix_rmask : bool

    Returns
    -------
    CopyNumArray
        One object summarizing the coverages of the input samples, including
        each bin's "average" coverage, "spread" of coverages, and GC content.
    """
    # Special parameters:
    # skip_low (when centering) = True for target; False for antitarget
    # do_edge  = as given for target; False for antitarget
    # do_rmask  = False for target; as given for antitarget
    ref_df, all_logr, all_depths = load_sample_block(
        filenames, fa_fname, is_haploid_x, sexes,
        True, fix_gc, fix_edge, False)
    if antitarget_fnames:
        # XXX TODO ensure ordering matches targets!
        #   argsort on both -> same?
        anti_ref_df, anti_logr, anti_depths = load_sample_block(
            antitarget_fnames, fa_fname, is_haploid_x, sexes,
            False, fix_gc, False, fix_rmask)
        ref_df = ref_df.append(anti_ref_df, ignore_index=True, sort=False)
        all_logr = np.hstack([all_logr, anti_logr])
        all_depths = np.hstack([all_depths, anti_depths])

    stats_all = summarize_info(all_logr, all_depths)
    ref_df = ref_df.assign(**stats_all)

    if do_cluster:
        # Get extra cols, concat axis=1 here (DATAFRAME v-concat)
        sample_ids = [core.fbase(f) for f in filenames]
        if len(sample_ids) != len(all_logr) - 1:
            raise ValueError("Expected %d target coverage files (.cnn), got %d"
                             % (len(all_logr) - 1, len(sample_ids)))
        clustered_cols = create_clusters(all_logr, min_cluster_size, sample_ids)
        if clustered_cols:
            try:
                ref_df = ref_df.assign(**clustered_cols)
            except ValueError as exc:
                print("Reference:", len(ref_df.index))
                for cl_key, cl_col in clustered_cols.items():
                    print(cl_key, ":", len(cl_col))
                raise exc
        else:
            print("** Why weren't there any clustered cols?")

    ref_cna = CNA(ref_df, meta_dict={'sample_id': 'reference'})
    # NB: Up to this point, target vs. antitarget bins haven't been row-sorted.
    # Holding off until here ensures the cluster-specific log2 columns are
    # sorted with the same row order as the rest of the CNA dataframe.
    ref_cna.sort()
    ref_cna.sort_columns()
    # TODO figure out centering
    #ref_probes.center_all(skip_low=True)  # <-- on each cluster col, too?
    #   or just subtract the same amount from the cluster calls after figuring
    #   out the main one?
    return ref_cna

def load_sample_block(filenames, fa_fname,
                      is_haploid_x, sexes, skip_low,
                      fix_gc, fix_edge, fix_rmask):
    r"""Load and summarize a pool of \*coverage.cnn files.

    Run separately for the on-target and (optional) antitarget bins.

    Returns
    -------
    ref_df : pandas.DataFrame
        All columns needed for the reference CNA object, including
        aggregate log2 and spread.
    all_logr : numpy.ndarray
        All sample log2 ratios, as a 2D matrix (rows=bins, columns=samples),
        to be used with do_cluster.
    """
    # Ensures samples' target and antitarget matrix columns are in the same
    # order, so they can be concatenated.
    # (We don't explicitly pair off each sample's target and antitarget .cnn
    # files; as long as the filename prefixes match, the resulting columns will
    # be consistent. Same is true for sample sex inference.)
    filenames = sorted(filenames, key=core.fbase)

    # Load coverage from target/antitarget files
    logging.info("Loading %s", filenames[0])
    cnarr1 = read_cna(filenames[0])
    if not len(cnarr1):
        # Just create an empty array with the right columns
        col_names = ['chromosome', 'start', 'end', 'gene', 'log2', 'depth']
        if 'gc' in cnarr1 or fa_fname:
            col_names.append('gc')
        if fa_fname:
            col_names.append('rmask')
        col_names.append('spread')
        empty_df = pd.DataFrame.from_records([], columns=col_names)
        empty_logr = np.array([[]] * (len(filenames) + 1))
        empty_dp = np.array([[]] * len(filenames))
        return empty_df, empty_logr, empty_dp

    # Calculate GC and RepeatMasker content for each probe's genomic region
    ref_columns = {
        'chromosome': cnarr1.chromosome,
        'start': cnarr1.start,
        'end': cnarr1.end,
        'gene': cnarr1['gene'],
    }
    if fa_fname and (fix_rmask or fix_gc):
        gc, rmask = get_fasta_stats(cnarr1, fa_fname)
        if fix_gc:
            ref_columns['gc'] = gc
        if fix_rmask:
            ref_columns['rmask'] = rmask
    elif 'gc' in cnarr1 and fix_gc:
        # Reuse .cnn GC values if they're already stored (via import-picard)
        gc = cnarr1['gc']
        ref_columns['gc'] = gc

    # Make the sex-chromosome coverages of male and female samples compatible
    is_chr_x = (cnarr1.chromosome == cnarr1._chr_x_label)
    is_chr_y = (cnarr1.chromosome == cnarr1._chr_y_label)
    ref_flat_logr = cnarr1.expect_flat_log2(is_haploid_x)
    ref_edge_bias = fix.get_edge_bias(cnarr1, params.INSERT_SIZE)
    # Pseudocount of 1 "flat" sample
    all_depths = [cnarr1['depth'] if 'depth' in cnarr1
                  else np.exp2(cnarr1['log2'])]
    all_logr = [
        ref_flat_logr,
        bias_correct_logr(cnarr1, ref_columns, ref_edge_bias,
                          ref_flat_logr, sexes, is_chr_x, is_chr_y,
                          fix_gc, fix_edge, fix_rmask, skip_low)]

    # Load only coverage depths from the remaining samples
    for fname in filenames[1:]:
        logging.info("Loading %s", fname)
        cnarrx = read_cna(fname)
        # Bin information should match across all files
        if not np.array_equal(
                cnarr1.data.loc[:, ('chromosome', 'start', 'end', 'gene')].values,
                cnarrx.data.loc[:, ('chromosome', 'start', 'end', 'gene')].values):
            raise RuntimeError("%s bins do not match those in %s"
                               % (fname, filenames[0]))
        all_depths.append(cnarrx['depth'] if 'depth' in cnarrx
                          else np.exp2(cnarrx['log2']))
        all_logr.append(
            bias_correct_logr(cnarrx, ref_columns, ref_edge_bias, ref_flat_logr,
                              sexes, is_chr_x, is_chr_y,
                              fix_gc, fix_edge, fix_rmask, skip_low))
    all_logr = np.vstack(all_logr)
    all_depths = np.vstack(all_depths)
    ref_df = pd.DataFrame.from_dict(ref_columns)
    return ref_df, all_logr, all_depths

def bias_correct_logr(cnarr, ref_columns, ref_edge_bias,
                      ref_flat_logr, sexes, is_chr_x, is_chr_y,
                      fix_gc, fix_edge, fix_rmask, skip_low):
    """Perform bias corrections on the sample."""
    cnarr.center_all(skip_low=skip_low)
    shift_sex_chroms(cnarr, sexes, ref_flat_logr, is_chr_x, is_chr_y)
    # Skip bias corrections if most bins have no coverage (e.g. user error)
    if (cnarr['log2'] > params.NULL_LOG2_COVERAGE - params.MIN_REF_COVERAGE
       ).sum() <= len(cnarr) // 2:
        logging.warning("WARNING: most bins have no or very low coverage; "
                        "check that the right BED file was used")
    else:
        if 'gc' in ref_columns and fix_gc:
            logging.info("Correcting for GC bias...")
            cnarr = fix.center_by_window(cnarr, .1, ref_columns['gc'])
        if 'rmask' in ref_columns and fix_rmask:
            logging.info("Correcting for RepeatMasker bias...")
            cnarr = fix.center_by_window(cnarr, .1, ref_columns['rmask'])
        if fix_edge:
            logging.info("Correcting for density bias...")
            cnarr = fix.center_by_window(cnarr, .1, ref_edge_bias)
    return cnarr['log2']

def shift_sex_chroms(cnarr, sexes, ref_flat_logr, is_chr_x, is_chr_y):
    """Shift sample X and Y chromosomes to match the reference sex.

    Reference values::

        XY: chrX -1, chrY -1
        XX: chrX 0, chrY -1

    Plan::

        chrX:
        xx sample, xx ref: 0    (from 0)
        xx sample, xy ref: -= 1 (from -1)
        xy sample, xx ref: += 1 (from 0)    +1
        xy sample, xy ref: 0    (from -1)   +1
        chrY:
        xx sample, xx ref: = -1 (from -1)
        xx sample, xy ref: = -1 (from -1)
        xy sample, xx ref: 0    (from -1)   +1
        xy sample, xy ref: 0    (from -1)   +1

    """
    is_xx = sexes.get(cnarr.sample_id)
    cnarr['log2'] += ref_flat_logr
    if is_xx:
        # chrX has same ploidy as autosomes; chrY is just unusable noise
        cnarr[is_chr_y, 'log2'] = -1.0  # np.nan is worse
    else:
        # 1/2 #copies of each sex chromosome
        cnarr[is_chr_x | is_chr_y, 'log2'] += 1.0

def summarize_info(all_logr, all_depths):
    """Average & spread of log2ratios and depths for a group of samples.

    Can apply to all samples, or a given cluster of samples.
    """
    logging.info("Calculating average bin coverages")
    cvg_centers = np.apply_along_axis(descriptives.biweight_location, 0,
                                      all_logr)
    depth_centers = np.apply_along_axis(descriptives.biweight_location, 0,
                                        all_depths)
    logging.info("Calculating bin spreads")
    spreads = np.array([descriptives.biweight_midvariance(a, initial=i)
                        for a, i in zip(all_logr.T, cvg_centers)])

    result = {
        'log2': cvg_centers,
        'depth': depth_centers,
        'spread': spreads,
    }
    # TODO center the resulting log2
    #ref_df = pd.DataFrame.from_dict(ref_columns)
    #ref_cna = CNA.from_columns(ref_columns, {'sample_id': "reference"})
    return result

def create_clusters(logr_matrix, min_cluster_size, sample_ids):
    """Extract and summarize clusters of samples in logr_matrix.

    1. Calculate correlation coefficients between all samples (columns).
    2. Cluster the correlation matrix.
    3. For each resulting sample cluster (down to a minimum size threshold),
       calculate the central log2 value for each bin, similar to the full pool.
       Also print the sample IDs in each cluster, if feasible.

    Also recalculate and store the 'spread' of each cluster, though this might
    not be necessary/good.

    Return a DataFrame of just the log2 values. Column names are ``log2_i``
    where i=1,2,... .
    """
    from .cluster import markov, kmeans
    # Drop the pseudocount sample
    logr_matrix = logr_matrix[1:, :]
    print("Clustering", len(logr_matrix), "samples...")
    #clusters = markov(logr_matrix)
    clusters = kmeans(logr_matrix)
    cluster_cols = {}
    sample_ids = np.array(sample_ids)  # For easy indexing
    for i, clust_idx in enumerate(clusters):
        i += 1
        #print(len(clust_idx), clust_idx)
        if len(clust_idx) < min_cluster_size:
            logging.info("Skipping cluster #%d, size %d < min. %d",
                         i, len(clust_idx), min_cluster_size)
            continue
        logging.info("Summarizing cluster #%d of %d samples",
                        i, len(clust_idx))
        # List which samples are in each cluster
        samples = sample_ids[clust_idx]
        logging.info("\n".join(["\t" + s for s in samples]))
        # Calculate each cluster's summary stats
        clust_matrix = logr_matrix[clust_idx, :]
        # XXX re-add the pseudocount sample to each cluster? need benchmark
        clust_info = summarize_info(clust_matrix, [])
        cluster_cols.update({
            'log2_%d' % i: clust_info['log2'],
            'spread_%d' % i: clust_info['spread'],
        })
    return cluster_cols

def warn_bad_bins(cnarr, max_name_width=50):
    """Warn about target bins where coverage is poor.

    Prints a formatted table to stderr.
    """
    bad_bins = cnarr[fix.mask_bad_bins(cnarr)]
    fg_index = ~bad_bins['gene'].isin(params.ANTITARGET_ALIASES)
    fg_bad_bins = bad_bins[fg_index]
    if len(fg_bad_bins) > 0:
        bad_pct = (100 * len(fg_bad_bins)
                   / sum(~cnarr['gene'].isin(params.ANTITARGET_ALIASES)))
        logging.info("Targets: %d (%s) bins failed filters "
                     "(log2 < %s, log2 > %s, spread > %s)",
                     len(fg_bad_bins),
                     "%.4f" % bad_pct + '%',
                     params.MIN_REF_COVERAGE,
                     -params.MIN_REF_COVERAGE,
                     params.MAX_REF_SPREAD)
        if len(fg_bad_bins) < 500:
            gene_cols = min(max_name_width, max(map(len, fg_bad_bins['gene'])))
            labels = fg_bad_bins.labels()
            chrom_cols = max(labels.apply(len))
            last_gene = None
            for label, probe in zip(labels, fg_bad_bins):
                if probe.gene == last_gene:
                    gene = '  "'
                else:
                    gene = probe.gene
                    last_gene = gene
                if len(gene) > max_name_width:
                    gene = gene[:max_name_width-3] + '...'
                if 'rmask' in cnarr:
                    logging.info("  %s  %s  log2=%.3f  spread=%.3f  rmask=%.3f",
                                 gene.ljust(gene_cols),
                                 label.ljust(chrom_cols),
                                 probe.log2, probe.spread, probe.rmask)
                else:
                    logging.info("  %s  %s  log2=%.3f  spread=%.3f",
                                 gene.ljust(gene_cols),
                                 label.ljust(chrom_cols),
                                 probe.log2, probe.spread)

    # Count the number of BG bins dropped, too (names are all "Antitarget")
    bg_bad_bins = bad_bins[~fg_index]
    if len(bg_bad_bins) > 0:
        bad_pct = (100 * len(bg_bad_bins)
                   / sum(cnarr['gene'].isin(params.ANTITARGET_ALIASES)))
        logging.info("Antitargets: %d (%s) bins failed filters",
                     len(bg_bad_bins), "%.4f" % bad_pct + '%')

def get_fasta_stats(cnarr, fa_fname):
    """Calculate GC and RepeatMasker content of each bin in the FASTA genome."""
    logging.info("Calculating GC and RepeatMasker content in %s ...", fa_fname)
    gc_rm_vals = [calculate_gc_lo(subseq)
                  for subseq in fasta_extract_regions(fa_fname, cnarr)]
    gc_vals, rm_vals = zip(*gc_rm_vals)
    return np.asfarray(gc_vals), np.asfarray(rm_vals)

def calculate_gc_lo(subseq):
    """Calculate the GC and lowercase (RepeatMasked) content of a string."""
    cnt_at_lo = subseq.count('a') + subseq.count('t')
    cnt_at_up = subseq.count('A') + subseq.count('T')
    cnt_gc_lo = subseq.count('g') + subseq.count('c')
    cnt_gc_up = subseq.count('G') + subseq.count('C')
    tot = float(cnt_gc_up + cnt_gc_lo + cnt_at_up + cnt_at_lo)
    if not tot:
        return 0.0, 0.0
    frac_gc = (cnt_gc_lo + cnt_gc_up) / tot
    frac_lo = (cnt_at_lo + cnt_gc_lo) / tot
    return frac_gc, frac_lo

def fasta_extract_regions(fa_fname, intervals):
    """Extract an iterable of regions from an indexed FASTA file.

    Input: FASTA file name; iterable of (seq_id, start, end) (1-based)
    Output: iterable of string sequences.
    """
    with pyfaidx.Fasta(fa_fname, as_raw=True) as fa_file:
        for chrom, subarr in intervals.by_chromosome():
            logging.info("Extracting sequences from chromosome %s", chrom)
            for _chrom, start, end in subarr.coords():
                yield fa_file[_chrom][int(start):int(end)]

def reference2regions(refarr):
    """Split reference into target and antitarget regions."""
    is_bg = (refarr['gene'].isin(params.ANTITARGET_ALIASES))
    regions = GA(refarr.data.loc[:, ('chromosome', 'start', 'end', 'gene')],
                 {'sample_id': 'reference'})
    targets = regions[~is_bg]
    antitargets = regions[is_bg]
    return targets, antitargets


"""
This is an implementation of paper
"Attention-based LSTM for Aspect-level Sentiment Classification" with Keras.
Based on dataset from "SemEval 2014 Task 4".
"""

import os
from time import time
# TODO, Here we need logger!

import numpy as np
from lxml import etree
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from keras.utils.np_utils import to_categorical
from keras.layers import Input, Embedding, LSTM, Dense
from keras.layers import RepeatVector, Dot, Concatenate, Reshape
from keras.activations import softmax
from keras.models import Model, load_model
from keras import regularizers, initializers, optimizers
from keras.layers import Lambda
import keras.backend as K

TEXT_KEY = 'text'
TERM_KEY = 'aspect_terms'
CATEGORY_KEY = 'aspect_categories'

I_TEXT, I_ASPECT, I_POLARITY = 0, 1, 2

# Correspond to settings in paper.
EMBEDDING_DIM = 300
ASPECT_EMBEDDING_DIM = 300
HIDDEN_LAYER_SIZE = 300

# Hyper-parameters for training.
L2_REGULARIZATION = 0.001
MOMENTUM = 0.9
LEARNING_RATE = 0.001
MINI_BATCH_SIZE = 25

RANDOM_UNIFORM = .01

POLARITY_TO_INDEX = {
    'positive': 0,
    'negative': 1,
    'neutral': 2,
    'conflict': 3
}

def extract_data(data_file='Restaurants_Train_v2.xml'):
    """
    Extract train data from xml file provided buy 'SemEval 2014 Task 4."
    :param file: XML file that contains training data.
    :return: A list of dictionaries of training data with TEXT_KEY, 'aspect
    terms' and 'aspect categories'.
    """
    tree = etree.parse(data_file)
    sents_root = tree.getroot()
    data = []

    def get_content(sent):
        """
        Get all contents from a single 'sentence node', including TEXT_KEY,
        values of 'aspect terms' and 'aspect categories'.
        :param sent: a single xml node of sentence.
        :type: _Element
        :return: A dictionary of contents.
        """
        content = {}
        # We assume that there is must a text node here.
        content[TEXT_KEY] = sent.xpath(TEXT_KEY)[0].text

        terms = sent.xpath('aspectTerms')
        if terms:
            # As there is only one element of 'aspectTerms'.
            # And we only need the first two values, 'aspect' and 'polarity'.
            content[TERM_KEY] = list(map(lambda term: term.values()[:2],
                                         terms[0].iterchildren()))
        else:
            pass

        categories = sent.xpath('aspectCategories')
        if categories:
            content[CATEGORY_KEY] = list(
                map(lambda category: category.values(),
                    categories[0].iterchildren()))
        else:
            pass

        return content

    for sent in sents_root.iterchildren():
        data.append(get_content(sent))

    return data

def check_absent(data):
    """
    Checking absent 'aspect terms' or 'aspect categories'.
    And check if there is sentence missing both 'terms' and 'categories'.
    :param data: dataset with all contents. And the max length of all sentence.
    :type: list of dictionary.
    :return: sentence indices that with absent terms, categories and flag of
    both missing as well as their count separately in tuple.
    :type: tuple of (list, list, boolean)
    """
    exist_both_missing = False
    term_absent_indices = []
    term_absent_cnt = 0
    category_absent_indices = []
    category_absent_cnt = 0
    max_len = 0

    for idx, sent in enumerate(data):
        max_len = max(len(sent[TEXT_KEY]), max_len)

        term_absent = TERM_KEY not in sent.keys()
        category_absent = CATEGORY_KEY not in sent.keys()
        if term_absent and category_absent:
            exist_both_missing = True

        if term_absent:
            term_absent_indices.append(idx)
            term_absent_cnt += 1

        if category_absent:
            category_absent_indices.append(idx)
            category_absent_cnt += 1

    return (term_absent_indices, term_absent_cnt,
            category_absent_indices, category_absent_cnt,
            exist_both_missing, max_len)

def combine_data(data, mess=True, replace_space=True, replace_space_char='_'):
    """
    If `mess` is True, means we would mess all data together.
    Combine text with all aspects related to it, both aspect
    terms and aspect categories. And mess them up.
    But if `mess` is False. we will combined TEXT_KEY and aspect separately
    with 'terms' or 'categories', and return them as tuple.
    And also return the max length of sentence per term or category
    if `mess` is True or separate max length if `mess` is False.
    :param data: all data with TEXT_KEY and lists of 'aspect terms' and
    'categories'.
    :return: all combined data or combined data with 'aspect terms' and
    'categories' separately along with their max length or in all.
    """
    term_data, category_data = [], []
    term_max_len, category_max_len = 0, 0

    # TODO, How do we treat multi-word token as aspect term?
    # 1. take whole as one token an replace space with other mask.
    # 2. split into multiple tokens and average all embeddings.
    # 3. only take one word into consideration.
    # Note for aspect terms, it could contains spaces in the word, so should
    # not use space to split tokenizer, and take all as one token.
    # And also, there are other special characters in the phrase, like '-'.
    # They should be keep.
    for sent in data:
        text = sent[TEXT_KEY]
        is_term_exist = TERM_KEY in sent.keys()
        is_category_exist = CATEGORY_KEY in sent.keys()
        if is_term_exist:
            term_max_len = max(term_max_len, len(sent[TEXT_KEY]))
            for term, polarity in sent[TERM_KEY]:
                if replace_space:
                    term = term.replace(' ', replace_space_char)
                term_data.append([text, term, polarity])
        if is_category_exist:
            category_max_len = max(category_max_len, len(sent[TEXT_KEY]))
            for category, polarity in sent[CATEGORY_KEY]:
                if replace_space:
                    category = category.replace(' ', replace_space_char)
                category_data.append([text, category, polarity])

    # print(len(term_data), len(category_data))

    if mess:
        max_len = max(term_max_len, category_max_len)
        term_data.extend(category_data)
        return term_data, max_len
    else:
        return (term_data, term_max_len), (category_data, category_max_len)

def convert_data(data, max_len=None, with_label=True, extra_data=False):
    """
    Convert data to tuples of (word_vectors, aspect_indices, polarity) to
    word indices sequences and labels to one hot. In order to lookup in
    embedding layer.
    And convert polarity to class identifier, as defined by default in
    polarity to index.

    NOTE: keep in mind to match label and 'text' and 'aspect'!

    :param data: List of data with element of (text, aspect, polarity).
    :param word_vectors: Word Vector lookup table.
    :param with_label: Whether it is training data with label or
        test/customized data without label.
    :return: Arrays contain (word vectors, aspect indices, polarity class
    index), and each of them is a numpy array, along with the word to index
    dictionary.
    :type: numpy array.
    """
    # Set indicator for 'text', 'aspect' and 'polarity(label)'.
    converted_data, lookups = [], []

    texts, aspects, labels = [], [], []
    # TODO, we should count max length here?!
    for d in data:
        texts.append(d[I_TEXT])
        aspects.append(d[I_ASPECT])
        if with_label:
            labels.append(d[I_POLARITY])

    def convert_to_indices(examples, max_len=None, need_tokenizer=False,
                           customized_filter='!"#$%&()*+,-./:;<=>?@[\\]^_`{|}~\t\n'):
        """
        Fit and convert word to indices sequences and word index lookup, and if
        needed, return tokenizer as well.
        :param examples: list of words or sentences.
        :param max_len: the max length of indices sequences.
        :param need_tokenizer: return tokenizer or not.
        :type: boolean
        :return: (indices sequence, word index lookup, <tokenizer>)
        :type: tuple
        """
        tokenizer = Tokenizer(filters=customized_filter)
        tokenizer.fit_on_texts(examples)
        seqs = tokenizer.texts_to_sequences(examples)
        word_idx = tokenizer.word_index
        # TODO, do we need to pad, if yes, 'pre' or 'post'?
        if max_len:
            seqs = pad_sequences(seqs, maxlen=max_len)

        if need_tokenizer:
            return seqs, word_idx, tokenizer
        else:
            return seqs, word_idx

    text_seqs, text_word_idx = convert_to_indices(texts, max_len)
    converted_data.append(np.asarray(text_seqs, dtype='int32'))
    lookups.append(text_word_idx)

    # For aspect term maybe we should not use tokenizer and filter.
    aspects_seqs, aspects_idx = convert_to_indices(
        aspects,
        # TODO, should use less filter.
        customized_filter='#$%&/:;<=>?@[\\]^`{|}~\t\n')
    converted_data.append(np.asarray(aspects_seqs, dtype='int32'))
    lookups.append(aspects_idx)

    if with_label:
        labels_seqs, labels_idx = convert_to_indices(labels)

        # Normalize label sequences as we only need '4' classes and do not need
        # extra class for 'other'.
        labels_arr = np.asarray(labels_seqs, dtype='int') - 1
        labels_one_hot = to_categorical(labels_arr)  # aspects_seqs,
        # [:, np.newaxis],
        converted_data.append(labels_one_hot)
        lookups.append(labels_idx)
        # print(aspects_seqs)

        #    # Preprocessing text without max number of words.
        #    text_tokenizer = Tokenizer()
        #    text_tokenizer.fit_on_texts(texts)
        #    text_seqs = text_tokenizer.texts_to_sequences(texts)
        #    text_word_idx = text_tokenizer.word_index
        #    # Just get indices of words, and does not categorize it as we won't
        #    # multiply one-hot vector in practice as it is computation costly.
        #    # Instead we just lookup with embedding layer.
        #    text_data = pad_sequences(text_seqs, maxlen=max_len)
        #
        #    # Preprocessing aspects.
        #    # The same as word in text, it will be lookup in embedding layer.
        #    aspects_tokenizer = Tokenizer()
        #    aspects_tokenizer.fit_on_texts(aspects)
        #    aspects_seqs = aspects_tokenizer.texts_to_sequences(aspects)
        #    aspects_idx = aspects_tokenizer.word_index
        #
        #    # Processing labels
        #    # Convert labels from words into indices and then to one-hot categorical
        #    # indices.
        #    labels_tokenizer = Tokenizer()
        #    labels_tokenizer.fit_on_texts(labels)
        #    labels_seqs = labels_tokenizer.texts_to_sequences(labels)
        #    labels_idx = labels_tokenizer.

    return converted_data, lookups

def load_w2v(idxes, emb_file, save_to_file=None):
    """
    Load pre-trained embedding and match words in training data to form a
    small set of word embedding matrix with OOV with all '0's.
    NOTE: Keras tokenizer.word_index start from 1, in order to use '0'
    padding in pad_sequence and mask_zero in embedding layer and following
    layer.
    :param idxes: the word loopup dictionary of word indices.
    :param emb_file: pre-trained embedding file.
    :return: word embedding matrix fit for the training data.
    """
    # Only need the lookup for 'text'.
    idx = idxes[I_TEXT]

    # Initial word embedding matrix with all '0's.
    # TODO, here we could set embedding dimesion automatically.
    emb_matrix = np.zeros((len(idx) + 1, EMBEDDING_DIM))

    # Timing it.
    start_time = time()
    with open(emb_file) as emb:
        for line in emb:
            pieces = line.strip().split()
            word, coef = pieces[0].strip(), pieces[1:]
            begin_idx = 0
            for elem_idx, elem in enumerate(coef):
                # In case there is space in the word,
                # continuously test if the string could be interpret as float,
                # if yes, it means this piece element is the beginning of the
                # coefficient and if no, then append to word as part of the
                # token.
                try:
                    # Test if an element in coefficient is an actual
                    # coefficient of a part of key token.
                    float(elem)
                    # Record begin index of actual coefficient.
                    begin_idx = elem_idx + 1
                    # Only break when we find the begin index of actual
                    # coefficient.
                    break
                except Exception as e:
                    word += elem
                    # print(e)
                    # TODO, we could record the trail and error in log.
                    # print("Filed to load record with word: '{}' and "
                    #       "coefficient: {}".format(word, coef))
                    # print(word)
            coef = np.asarray(pieces[begin_idx:], dtype=np.float32)

            if word in idx.keys():
                # Lookup the indices(index) of word and set the corresponding
                # vector to the one in pre-trained embedding matrix.
                emb_matrix[idx[word]] = coef
        print('Loaded word embedding matrix within {}'.format(
            time() - start_time))

    # Save loaded subset of word embedding into files.
    if save_to_file:
        np.save(save_to_file, emb_matrix)

    return emb_matrix

def build_net(data, max_len, w2is, atae=True, extra_outputs=True,
              emb_mtrx_file=None, save_to_file=None):
    """
    Build ATAE-LSTM mentioned in paper 'Attention-based LSTM for Aspect-level
    Sentiment Classification', with uniform randomly initialized aspect
    embedding and word embedding subset according training data and given
    pre-trained embedding file.
    Adapt 'inter' attention before do multiple classes classification by
    softmax, which introduce aspect-level attention as part of the encoding
    of source sentence.`

    :param data: Indices of training data including (sentences, aspect,
        polarity(one-hot label))
    :param max_len: the max length of sentence as it has been padding with
    '0's and need to set for the input shape with mini-batch.
    :param w2is: Index lookup table of components above.
    :param atae: If 'False' then only use 'AE'.
    :param extra_outputs: return extra outputs like attention weights,
        aspect embeddings or so.
    :param emb_mtrx_file: Pre-saved embedding matrix corresponding to
        training data and given pre-trained embedding. If 'None' is set,
        then reload from embedding file.
    :param save_to_file: File path to save model, if 'None' is set, then its
    a one way training.
    :return: Training loss and accuracy for all classes?
    """
    # TODO, max length should be fixed.
    sents, aspects, labels = data
    sents_idx, aspects_idx, _ = w2is
    emb_mtrx = np.load(emb_mtrx_file)

    # Input of sentences.
    sents_tensor_input = Input(shape=(sents.shape[1],), dtype='int32')
    # Do not retrain embedding of sentences.
    sents_tensor = Embedding(len(sents_idx) + 1,
                             # EMBEDDING_DIM
                             emb_mtrx.shape[1],
                             weights=[emb_mtrx],
                             input_length=max_len,
                             trainable=False)(sents_tensor_input)

    # Input of aspect
    # As we use ATAE-LSTM, aspect embedding need to be concated to each time
    # steps in sentences.
    # Aspect is a single index of integer.
    aspects_tensor_input = Input(shape=(1,), dtype='int32')
    # Randomly initialize aspect embedding.
    aspects_emb_initializer = initializers.RandomUniform(minval=-RANDOM_UNIFORM,
                                                         maxval=RANDOM_UNIFORM)
    aspects_emb_layer = Embedding(len(aspects_idx) + 1,
                                  ASPECT_EMBEDDING_DIM,
                                  embeddings_initializer=aspects_emb_initializer,
                                  trainable=True,
                                  name='asp_emb_layer')
    # In order to get embedding weights.
    # aspects_emb_matrix = Lambda(lambda x: x, name='asp_emb_weight')(
    #     aspects_emb_layer.weights)
    aspects_emb = aspects_emb_layer(aspects_tensor_input)
    # Here, before repeat we need reshape aspect_tensor act as 'squeeze' with
    # the dimension with '1', say Reshape((10, ), input_shape=(1, 10))(...)
    # then got keras tensor with shape of (10,), which will then feed into
    # `RepeatVector`.
    aspects_tensor = Reshape((ASPECT_EMBEDDING_DIM,))(aspects_emb)
    # Repeat aspects tensor in order to correspond to the time step of
    # sentences, with shape of (max_len, ASPECT_EMBEDDNING_DIM).
    # TODO, could use Timedistributed?
    aspects_tensor = RepeatVector(max_len)(aspects_tensor)

    lstm_input = Concatenate()([sents_tensor, aspects_tensor])

    if atae:
        lstm_output = LSTM(HIDDEN_LAYER_SIZE, return_sequences=True)(lstm_input)
        # Attention with concatenation of sequential output of LSTM and
        # aspect embedding.
        attention_input = Concatenate()([lstm_output, aspects_tensor])
        attention_score = Dense(EMBEDDING_DIM + ASPECT_EMBEDDING_DIM,
                                use_bias=False,
                                name='attention_score_1')(attention_input)
        # We need an extra `Dense/Activation` layer here for axis related
        # softmax with should be align on time step instead the last axis.
        attention_weight = Dense(1, use_bias=False,
                                 name='attention_score_2')(attention_score)
        attention_weight = Lambda(lambda x: softmax(x, axis=1))(
            attention_weight, name='attention_weights')

        # permuted_weight = Permute((2, 1))(attention_weight)
        # attention_represent = Multiply(name='r')([lstm_output, permuted_weight])
        # attention_represent = Multiply(name='r')([lstm_output, attention_weight])
        attention_represent = Dot(axes=1, name='r')([lstm_output,
                                                     attention_weight])
        attention_represent = Reshape((EMBEDDING_DIM,))(attention_represent)

        last_hidden = Lambda(lambda tensor: tensor[:, -1, :])(lstm_output)
        final_represent = Concatenate(name='final_concatenate')([
            attention_represent, last_hidden])
        final_represent = Dense(EMBEDDING_DIM, activation='tanh',
                                use_bias=False, name='final_representation')(
            final_represent)
        model_output = Dense(labels.shape[1],
                             activation='softmax',
                             activity_regularizer=regularizers.l2(
                                 L2_REGULARIZATION),
                             name='ATAE_LSTM_output')(final_represent)

        # outs = [model_output]
        # if extra_outputs:
        #     outs.append(attention_weight)
        #     TODO, get from model outside
        # outs.append(aspects_emb_matrix)
        # print(outs)
    else:
        lstm_output = LSTM(HIDDEN_LAYER_SIZE,
                           return_sequences=False)(lstm_input)
        model_output = Dense(labels.shape[1],
                             activation='softmax',
                             name='Simple_AE_LSTM_ouptut')(lstm_output)
        # outs = [model_output]

    model = Model(inputs=[sents_tensor_input,
                          aspects_tensor_input],
                  outputs=model_output)

    if save_to_file:
        model.save(save_to_file)

    return model

def train(data, model, model_optimizer=None, metrics=None, valid_ratio=0.1,
          epoch=10, mini_batch=25, save_to_file=None):
    """
    :param data: Training data in tuples of lists with form of (sentences,
        aspect word, polarity).
    :param model: Predefined model generated by `build_net`, if None,
        then if will be build with default values.
    :param optimizer: Optimizer used to train/compile model. Default is
        Adagrad with learning rate as '0.001'.
    :param metrics: Metrics are interested in list. If not set then default
        is ['accuracy']
    :return: None
    """

    if not model and not data:
        print('Please passed in data and model!')
        return

    if not metrics:
        metrics = ['accuracy']

    if not model_optimizer:
        model_optimizer = optimizers.Adagrad(lr=0.001)

    print("Training Model ...")
    print(model.summary())
    # print('\t\twith data as')
    # print('\t\t{}'.format(check_absent(data)))
    print('\t\twith hyper-parametes as')
    print('\t\t\tMini-Batch : {}'.format(mini_batch))
    print('\t\t\tEpoch : {}'.format(epoch))

    model.compile(model_optimizer, 'categorical_crossentropy', metrics=metrics)
    model.fit([seqs_data[I_TEXT], seqs_data[I_ASPECT]], seqs_data[I_POLARITY],
              mini_batch, epochs=epoch, validation_split=valid_ratio)

    if save_to_file:
        model.save(save_to_file)

def train_dev_split(data, ratio=0.8, seed=42):
    """
    Function to split train and dev set with given ratio.
    :param data: whole dataset.
    :param ratio: percentage that training data occupied.
    :return: tuple of list of (training, dev), and each of them should be
        formed as (sentences, aspect word, polarity)
    """
    np.random.seed(42)
    sents, aspects, labels = data[I_TEXT], data[I_ASPECT], data[I_POLARITY]
    idx = np.arange(sents.shape[0])
    np.random.shuffle(idx)

    sents = sents[idx]
    aspects = aspects[idx]
    labels = labels[idx]

    # Calculate split boundary.
    bnd = int(len(idx) * ratio)

    train_set = [sents[:bnd], aspects[:bnd], labels[:bnd]]
    dev_set = [sents[bnd:], aspects[bnd:], labels[bnd:]]

    return train_set, dev_set

def predict(data, lookup, max_len, model=None, save_to_file=None,
            extra_output=True):
    """
    Predict with given data and model or load model from saved pre-trained
    model in file.
    :param data: data in tuple or list (sentence, aspect)
    :param w2is: index to lookup for predictions.
    :param max_len: length to padding to.
    :param model: pre-trained model, if not set loaded from file,
        and if file for model is also not set, return with error.
    :param save_to_file: file saved with model.
    :return: prediction
    """
    # Omit word index lookups.
    converted_data, _ = convert_data(data, max_len, with_label=False)
    # print(converted_data)
    if not model:
        if save_to_file:
            model = load_model(save_to_file,
                               custom_objects={'softmax': softmax})
        else:
            # TODO, should raise exception?
            raise ValueError('Please pass in model instance or '
                             'the path of file model saved to.')
    pred_vec = model.predict([converted_data[I_TEXT],
                              converted_data[I_ASPECT]])
    pred_idx = np.argmax(pred_vec, axis=1)
    func_get_label = np.vectorize(lambda p: lookup.get(p))
    # print(pred_idx, func_get_label(pred_idx), lookup.get(0))
    # Need to add '1' for keras labels start from '0'.
    pred = func_get_label(pred_idx + 1)

    # if extra_output:
    #     model.layers

    return pred

def get_layer(model, layer_name):
    """
    Get layer from model by name or index.
    :param layer_name: the name or index of layer.
    :return: layer instance extract from model.
    """
    if isinstance(layer_name, int):
        return model.layers[layer_name]
    elif isinstance(layer_name, str):
        return model.get_layer(layer_name)
    else:
        raise ValueError('The layer name should only be `int` or `str`.')

def get_aspect_embeddings(model, layer_name, save_to_file=None):
    """
    Get aspect embedding from specific layer with given name.
    :param model: the pre-trained model, if not set, reload form saved model
        file. If it also failed to load model from file, 'ValueError' will be thrown.
    :param layer_name: the name or index of embedding layer, or ValueError
        will be thrown.
    :param save_to_file: file saved pre-trained model, load model if model is 'None'.
    :return: tensor of apsect embeddings.
    """
    if not model:
        if not save_to_file:
            raise ValueError('No model found from parameter or file!')
        else:
            model = load_model(save_to_file)
    # Get embeddings of aspect words.
    emb_layer = get_layer(model, layer_name)
    return K.eval(emb_layer.embeddings)

def get_attention_weighs(data, att_layer_name, input_layers_names: list,
                         model=None, save_to_file=None):
    """
    Get attention weights(intermediate) from specific layer with given layer
        name and input layers.
    :param data: data to attendant to.
    :param model: the pre-trained model, if not set, reload form saved model
        file. If it also failed to load model from file, 'ValueError' will be thrown.
    :param att_layer_name: the name or index of embedding layer, or ValueError
        will be thrown.
    :param input_layers: the name or index list of all input layer in order.
    :param save_to_file: file saved pre-trained model, load model if model is 'None'.
    :return: tensor of attention indices.
    """
    if not model:
        if not save_to_file:
            raise ValueError('No model found from parameter or file!')
        else:
            model = load_model(save_to_file,
                               custom_objects={'softmax': softmax})

    # Must be sure input layers are in order.
    att_layer = get_layer(model, att_layer_name)

    input_layers = []
    for layer_name in input_layers_names:
        layer = get_layer(model, layer_name)
        if layer:
            input_layers.append(layer.input)

    get_attention_weights = K.function(input_layers, [att_layer.output])
    weights = get_attention_weights([data[I_TEXT], data[I_ASPECT]])[0]
    # print(weights.shape)

    return weights

def plot_attention_weight(weights, focus_len):
    """
    Plot attention weights within the focus length.
    :param weights: attention weights.
    :param focus_len: the length to focus to, usually the length of sentences.
    :return: None
    """
    # score_file = os.path.join(RAW_DATA_FILE_BASE, 'intermeidate_score')
    # np.save(score_file, weights)

    # score_input = Input(shape=(term_max_len, 600))
    # get_weights = Dense(1, use_bias=False)(score_input)
    # get_weights = Activation('softmax', axis=1)(get_weights)
    # get_weights = Lambda(lambda x: tf.nn.softmax())
    # from keras.activations import softmax

    # # # get_weights = Lambda(lambda x: softmax(x, axis=1))(get_weights)
    # # # score_model = Model(score_input, get_weights)
    # # # print(score_model.summary())
    # #
    # # score_model.compile(optimizer='adam', loss='categorical_crossentropy')
    # weight_result = score_model.predict(weights)
    # print(weight_result[0].shape)

    # begin_idx = len(converted_data[I_TEXT][0])
    # print(begin_idx)
    import matplotlib.pyplot as plt

    # hist, bins = np.histogram(weight_result[0].reshape((1, -1)))
    # We have to remember the length of input sentences in order to align the
    # attention weights.
    # plt.imshow(weight_result[0][-20:].reshape((1, -1)), cmap="plasma",
    #            aspect="auto", extent=[0, 20, 0, 1])

    # TODO, Here is 'pre pad', so its '-focus_len' for the actual token.
    attentions = weights.reshape((1, -1))[:, -focus_len:]
    print(attentions.shape)
    plt.imshow(attentions, cmap='plasma',
               aspect='auto', extent=[0, focus_len, 0, 1])
    # plt.grid(True)
    plt.colorbar()
    plt.show()

if __name__ == '__main__':
    RAW_DATA_FILE_BASE = '/Users/jiazhen/datasets/SemEval' \
                         '/SemEval_2014_task4/ABSA_v2'
    RES_RAW_DATA_FILE = os.path.join(RAW_DATA_FILE_BASE,
                                     'Restaurants_Train_v2.xml')
    LAP_RAW_DATA_FILE = os.path.join(RAW_DATA_FILE_BASE, 'Laptop_Train_v2.xml')

    WORD_EMB_BASE = '/Users/jiazhen/datasets'
    WORD_EMB_FILE = os.path.join(WORD_EMB_BASE, 'glove.840B.300d.txt')

    SAVED_EMB_FILE = os.path.join(RAW_DATA_FILE_BASE, 'glove_res_emb.npy')
    SAVED_MDL_FILE = os.path.join(RAW_DATA_FILE_BASE, 'atae_model.keras')

    res_data = extract_data(RES_RAW_DATA_FILE)

    # print(res_data[7])
    check_absent(res_data)
    (term_data, term_max_len), _ = combine_data(res_data, mess=False)
    # print(term_data[7])

    # No padding here according to the paper.
    # Need padding for mini-batch.
    seqs_data, w2is = convert_data(term_data, max_len=term_max_len)

    # emb_matrix = load_w2v(w2is, WORD_EMB_FILE, SAVED_EMB_FILE)
    # print(emb_matrix[1])

    # print(len(seqs_data))
    # print(seqs_data[0].shape, seqs_data[1].shape, seqs_data[2].shape)
    # print(seqs_data[1])
    # for i, d in enumerate(seqs_data[1]):
    #     if len(d) > 1:
    #         print(i, d)
    #         print(term_data[i][I_ASPECT])
    # print('raw data', res_data[92]['aspect_terms'])
    # print(type(seqs_data[1][0][0]))
    # print(type(seqs_data[2][0][0]))
    # print(w2is[0])

    # reloaded_emb = np.load(SAVED_EMB_FILE)
    # print(reloaded_emb[1])

    # Train model.
    # model = build_net(seqs_data, term_max_len, w2is,
    #                   atae=True, extra_outputs=True,
    #                   emb_mtrx_file=SAVED_EMB_FILE,
    #                   save_to_file=SAVED_MDL_FILE + '2')
    # train(seqs_data, model, epoch=3)

    label_lookup = {idx: polarity
                    for polarity, idx in w2is[I_POLARITY].items()}
    # print(label_lookup)
    customized_data = [['The food is really delicious but '
                        'I hate the service', 'food'],
                       ['The food is really delicious but '
                        'I hate the service', 'serivce'],
                       ['I have to say there is no on could be faster than '
                        'him, but he need to take care of his bad motion as '
                        'a bar attendant, which will impact his serivce.',
                        'serivce']]
    pred = predict(customized_data, label_lookup, term_max_len,
                   save_to_file=SAVED_MDL_FILE + '2')
    print(pred)

    # Get attention weights for sentences.
    converted_data, _ = convert_data(customized_data,
                                     term_max_len,
                                     with_label=False)

    weights = get_attention_weighs(converted_data,
                                   att_layer_name='attention_weight',
                                   # att_layer_name='attention_weights',
                                   input_layers_names=[2, 0],
                                   save_to_file=SAVED_MDL_FILE + '2')
    # print(weights[0])
    print(len(customized_data[0][I_TEXT].split()))
    focus_len = len(customized_ata[0][I_TEXT].split())
    plot_attention_weight(weights[0], focus_len=focus_len)
    # for weight in weights:
    #     print(weight.shape)

    # TODO, Use gemsim to visualize aspect word embeddings.


import os
import torch
import argparse
import torch.optim as optim
from data import get_training_set
from network import Vgg19, Encoder, Img_decoder_v3, discriminator_v2, Inverse
from libs import MulLayer
import torchvision.utils as vutils
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import numpy as np
from criterion import LossCriterion_v2, LapLoss, TV, mean_variance_norm
from models import calc_kl, reparameterize
from lpips import lpips

parser = argparse.ArgumentParser()

parser.add_argument('--pretrained', type=bool, default=True)
parser.add_argument('--enc_dir', type=str, default='models/enc.pth')
parser.add_argument('--dec_dir', type=str, default='models/dec.pth')
parser.add_argument('--inv_dir', type=str, default='models/inv.pth')
parser.add_argument("--stylePath", default="data/sketch",
                    help='path to wikiArt dataset')
parser.add_argument("--contentPath", default="data/coco",
                    help='path to MSCOCO dataset')
parser.add_argument('--data_augmentation', type=bool, default=True)
parser.add_argument('--fineSize', type=int, default=256, help='crop image size')
parser.add_argument('--crop', type=bool, default=True, help='crop training images')
parser.add_argument('--threads', type=int, default=4, help='number of threads for data loader to use')
parser.add_argument("--outf", default="output/", help='folder to output images and model checkpoints')
parser.add_argument("--content_weight", type=float, default=0.1, help='content loss weight')
parser.add_argument("--style_weight", type=float, default=0.4, help='style loss weight, 0.02 for origin')
parser.add_argument("--batchSize", type=int, default=4, help='batch size')
parser.add_argument("--lr", type=float, default=2e-4, help='learning rate')
parser.add_argument("--gpu_id", type=int, default=0, help='which gpu to use')
parser.add_argument("--save_interval", type=int, default=5000, help='checkpoint save interval')
parser.add_argument("--layer", default="r41", help='which features to transfer, either r31 or r41')
parser.add_argument('--nEpochs', type=int, default=100, help='number of epochs to train for')
parser.add_argument('--snapshots', type=int, default=1, help='Snapshots')
parser.add_argument('--start_iter', type=int, default=1, help='Starting Epoch')

################# PREPARATIONS #################
opt = parser.parse_args()

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

os.makedirs(opt.outf, exist_ok=True)
cudnn.benchmark = True

def print_network(net):
    num_params = 0
    for param in net.parameters():
        num_params += param.numel()
    # print(net)
    print('Total number of parameters: %d' % num_params)

################# DATA #################
print('===> Loading datasets')
train_set = get_training_set(opt.contentPath, opt.stylePath)
training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True)

################# MODEL #################
enc = Encoder()
dec = Img_decoder_v3()
inv = Inverse()
VGG = Vgg19()
D = discriminator_v2(device, num_channels=1, base_filter=64)
criterion = LossCriterion_v2(opt.style_weight, opt.content_weight, device=device)

print('---------- encoder architecture -------------')
print_network(enc)

if opt.pretrained:
    if os.path.exists(opt.dec_dir):
        pretrained_dict = torch.load(opt.dec_dir, map_location=lambda storage, loc: storage)
        model_dict = dec.state_dict()
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
        model_dict.update(pretrained_dict)
        dec.load_state_dict(model_dict)
        # dec.load_state_dict(torch.load(opt.dec_dir))
        print('pretrained Decoder model is loaded!')
    if os.path.exists(opt.enc_dir):
        pretrained_dict = torch.load(opt.enc_dir, map_location=lambda storage, loc: storage)
        model_dict = enc.state_dict()
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
        model_dict.update(pretrained_dict)
        enc.load_state_dict(model_dict)
        # dec.load_state_dict(torch.load(opt.dec_dir))
        print('pretrained Encoder model is loaded!')
    if os.path.exists(opt.inv_dir):
        pretrained_dict = torch.load(opt.inv_dir, map_location=lambda storage, loc: storage)
        model_dict = inv.state_dict()
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
        model_dict.update(pretrained_dict)
        inv.load_state_dict(model_dict)
        # dec.load_state_dict(torch.load(opt.dec_dir))
        print('pretrained INverser model is loaded!')

# for param in inv.parameters():
#     param.requires_grad = False
################# LOSS & OPTIMIZER #################
L1_criterion_avg = torch.nn.L1Loss(size_average=True)
BCE_criterion = torch.nn.BCELoss()
Lap_criterion = LapLoss(device=device, max_levels=5)
optimizer_G = optim.AdamW(list(enc.parameters()) + list(dec.parameters()) + list(inv.parameters()), lr=opt.lr)
optimizer_D = optim.AdamW(D.parameters(), lr=opt.lr)

################# GPU  #################
# vgg = torch.nn.DataParallel(vgg)
# dec = torch.nn.DataParallel(dec)
# vgg5 = torch.nn.DataParallel(vgg5)
# matrix = torch.nn.DataParallel(matrix)
# OF_net = torch.nn.DataParallel(OF_net)

D.to(device)
enc.to(device)
dec.to(device)
inv.to(device)
VGG.to(device)
L1_criterion_avg.to(device)
BCE_criterion.to(device)
Lap_criterion.to(device)

################# TRAINING #################

def train(epoch):
    epoch_loss = 0
    enc.train()
    dec.train()
    D.train()
    inv.train()

    for iteration, batch in enumerate(training_data_loader, 1):
        img, edge, edge_gt, ref = batch[0], batch[1], batch[2], batch[3]
        img = img.to(device)
        edge = edge.to(device)
        edge_gt = edge_gt.to(device)
        ref = ref.to(device)

        b, c, h, w = img.shape
        real_label = torch.ones((b, 336)).to(device)
        fake_label = torch.zeros((b, 336)).to(device)

        img_gray = img[:, 0:1, :, :] * 0.299 + img[:, 1:2, :, :] * 0.587 + img[:, 2:3, :, :] * 0.114
        edge_gt = edge_gt[:, 0:1, :, :] * 0.299 + edge_gt[:, 1:2, :, :] * 0.587 + edge_gt[:, 2:3, :, :] * 0.114
        edge = (edge_gt + edge) / 2.0
        # edge = edge_gt * edge
        # edge = edge.clamp(0, 1)
        # img = img.repeat(1, 3, 1, 1)
        # forward
        for param in D.parameters():
            param.requires_grad = False

        optimizer_G.zero_grad()

        feat = enc(img)
        output = dec(feat)
        img_predict = inv(output)
        cF2 = VGG(img_predict)
        cF1 = VGG(img)
        sF = VGG(ref)
        tF = VGG(output.repeat(1, 3, 1, 1))

        # VGG feature
        Laploss = L1_criterion_avg(output, img_gray)
        Imgloss = Lap_criterion(img_predict.repeat(1, 3, 1, 1), img_gray.repeat(1, 3, 1, 1))
        Edgeloss = L1_criterion_avg(output, edge)
        loss_1, styleLoss, contentLoss1= criterion(tF, sF, cF1)
        _, _, contentLoss2 = criterion(cF2, sF, cF1)

        D_fake_feat, D_fake_decision = D(output)

        GAN_loss = L1_criterion_avg(D_fake_decision, real_label)

        G_loss = 5 * GAN_loss + 5 * Laploss + 100 * Imgloss + 100 * Edgeloss + loss_1 + 0.1 * contentLoss2

        # backward & optimization
        G_loss.backward()
        optimizer_G.step()

        # Reset gradient
        for p in D.parameters():
            p.requires_grad = True

        optimizer_D.zero_grad()

        ref = ref[:, 0:1, :, :] * 0.299 + ref[:, 1:2, :, :] * 0.587 + ref[:, 2:3, :, :] * 0.114
        _, D_real_decision = D(ref)
        _, D_fake_decision = D(output.detach())

        real = real_label * np.random.uniform(0.7, 1.2)
        fake = fake_label + np.random.uniform(0.0, 0.3)

        D_loss = (L1_criterion_avg(D_real_decision, real)
                    + L1_criterion_avg(D_fake_decision, fake)) / 2.0

        # Back propagation
        D_loss.backward()
        optimizer_D.step()

        print("===> VAE Epoch[{}]({}/{}): G_loss: {:.4f} || D_loss: {:.4f} || "
              "Imgloss: {:.4f} || Laploss: {:.4f} || Edgeloss: {:.4f} || contentloss1: {:.4f} || "
              "contentloss2: {:.4f} || styleloss: {:.4f} ||"
              "GAN_loss: {:.4f}".format(epoch, iteration,
                                                            len(training_data_loader),
                                                            G_loss.data, D_loss.data,
                                                            Imgloss.data, Laploss.data, Edgeloss.data,
                                                            contentLoss1.data, contentLoss2.data, styleLoss.data,
                                                            GAN_loss.data))

        if iteration % 500 == 0:
            img = img.clamp(0, 1).cpu().data
            edge = edge.repeat(1, 3, 1, 1).clamp(0, 1).cpu().data
            ref = ref.repeat(1, 3, 1, 1).clamp(0, 1).cpu().data
            output = output.repeat(1, 3, 1, 1).clamp(0, 1).cpu().data
            img_predict = img_predict.repeat(1, 3, 1, 1).clamp(0, 1).cpu().data
            concat = torch.cat((img, ref, edge, output, img_predict), dim=0)
            vutils.save_image(concat, '%s/%d_%d.png' % (opt.outf, epoch, iteration), normalize=True,
                              scale_each=True, nrow=img.shape[0])

            torch.save(enc.state_dict(), 'output/enc_iter_%d_epoch_%d.pth' % (iteration, epoch))
            torch.save(dec.state_dict(), 'output/dec_iter_%d_epoch_%d.pth' % (iteration, epoch))
            torch.save(inv.state_dict(), 'output/inv_iter_%d_epoch_%d.pth' % (iteration, epoch))

    return img, output

for epoch in range(opt.start_iter, opt.nEpochs + 1):
    img, output = train(epoch)
    # content = torch.cat((img1, img2), dim=3)
    # style = style.repeat(1, 1, 1, 2)
    # transfer = torch.cat((transfer1, transfer2), dim=3)
    # learning rate is decayed by a factor of 10 every half of total epochs


import argparse
import math
import numpy as np
import tensorflow as tf
import socket
import importlib
import os
import sys
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
from pointnetvlad_cls import *
from loading_pointclouds_ford import *
from sklearn.neighbors import NearestNeighbors
from sklearn.neighbors import KDTree
from tqdm import tqdm
import time

#params
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]')
parser.add_argument('--positives_per_query', type=int, default=4, help='Number of potential positives in each training tuple [default: 2]')
parser.add_argument('--negatives_per_query', type=int, default=12, help='Number of definite negatives in each training tuple [default: 20]')
parser.add_argument('--batch_num_queries', type=int, default=3, help='Batch Size during training [default: 1]')
parser.add_argument('--dimension', type=int, default=256)
parser.add_argument('--decay_step', type=int, default=200000, help='Decay step for lr decay [default: 200000]')
parser.add_argument('--decay_rate', type=float, default=0.7, help='Decay rate for lr decay [default: 0.8]')
FLAGS = parser.parse_args()

#BATCH_SIZE = FLAGS.batch_size
BATCH_NUM_QUERIES = FLAGS.batch_num_queries
EVAL_BATCH_SIZE = 1
NUM_POINTS = 4096
POSITIVES_PER_QUERY= FLAGS.positives_per_query
NEGATIVES_PER_QUERY= FLAGS.negatives_per_query
GPU_INDEX = FLAGS.gpu
DECAY_STEP = FLAGS.decay_step
DECAY_RATE = FLAGS.decay_rate

# LOG_DIR = '../models/refine'
LOG_DIR = 'log_fold07/' # todo
RESULTS_FOLDER= "KITTI_Ford"
# RESULTS_FOLDER="pretrained_results/"

if not os.path.exists(RESULTS_FOLDER): os.mkdir(RESULTS_FOLDER)

OUT_FEATURE_FOLDER = os.path.join(RESULTS_FOLDER, "feature_database")    # todo
if not os.path.exists(OUT_FEATURE_FOLDER): os.makedirs(OUT_FEATURE_FOLDER)
# DATABASE_FILE= 'generating_queries/oxford_evaluation_database.pickle'
# QUERY_FILE= 'generating_queries/oxford_evaluation_query.pickle'
KITTI_submap_dir = "/media/data/Ford/submap"

output_file= RESULTS_FOLDER +'results.txt'
# model_file= "model_refine.ckpt"
model_file= "model.ckpt"

# DATABASE_SETS= get_sets_dict(DATABASE_FILE)
# QUERY_SETS= get_sets_dict(QUERY_FILE)

# global DATABASE_VECTORS
# DATABASE_VECTORS=[]
#
# global QUERY_VECTORS
# QUERY_VECTORS=[]

BN_INIT_DECAY = 0.5
BN_DECAY_DECAY_RATE = 0.5
BN_DECAY_DECAY_STEP = float(DECAY_STEP)
BN_DECAY_CLIP = 0.99

def get_bn_decay(batch):
    bn_momentum = tf.train.exponential_decay(
                      BN_INIT_DECAY,
                      batch*BATCH_NUM_QUERIES,
                      BN_DECAY_DECAY_STEP,
                      BN_DECAY_DECAY_RATE,
                      staircase=True)
    bn_decay = tf.minimum(BN_DECAY_CLIP, 1 - bn_momentum)
    return bn_decay

def evaluate():
    # global DATABASE_VECTORS
    # global QUERY_VECTORS

    with tf.Graph().as_default():
        with tf.device('/gpu:'+str(GPU_INDEX)):
            print("In Graph")
            query= placeholder_inputs(BATCH_NUM_QUERIES, 1, NUM_POINTS)
            positives= placeholder_inputs(BATCH_NUM_QUERIES, POSITIVES_PER_QUERY, NUM_POINTS)
            negatives= placeholder_inputs(BATCH_NUM_QUERIES, NEGATIVES_PER_QUERY, NUM_POINTS)
            eval_queries= placeholder_inputs(EVAL_BATCH_SIZE, 1, NUM_POINTS)

            is_training_pl = tf.placeholder(tf.bool, shape=())
            print(is_training_pl)

            batch = tf.Variable(0)
            bn_decay = get_bn_decay(batch)

            with tf.variable_scope("query_triplets") as scope:
                vecs= tf.concat([query, positives, negatives],1)
                print(vecs)
                out_vecs= forward(vecs, is_training_pl, bn_decay=bn_decay)
                print(out_vecs)
                q_vec, pos_vecs, neg_vecs= tf.split(out_vecs, [1,POSITIVES_PER_QUERY,NEGATIVES_PER_QUERY],1)
                print(q_vec)
                print(pos_vecs)
                print(neg_vecs)

            saver = tf.train.Saver()

        # Create a session
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.95)
        config = tf.ConfigProto(gpu_options=gpu_options)
        config.gpu_options.allow_growth = True
        config.allow_soft_placement = True
        config.log_device_placement = False
        sess = tf.Session(config=config)

        saver.restore(sess, os.path.join(LOG_DIR, model_file))
        print("Model restored.")

        ops = {'query': query,
               'positives': positives,
               'negatives': negatives,
               'is_training_pl': is_training_pl,
               'eval_queries': eval_queries,
               'q_vec':q_vec,
               'pos_vecs': pos_vecs,
               'neg_vecs': neg_vecs}
        # recall= np.zeros(25)
        # count=0
        # similarity=[]
        # one_percent_recall=[]
        # sequences = ["00", "02", "05", "08", "06", "07"]
        sequences = ["01","02"]  # todo
        for sq in sequences:
            sq_dir = os.path.join(KITTI_submap_dir, sq)
            t1 = time.time()
            feature_db = get_latent_vectors(sess, ops, sq_dir)
            t2 = time.time()
            print("time: ", t2 - t1)
            feature_db_name = os.path.join(OUT_FEATURE_FOLDER, sq + "_PV_" + sq + ".npy")
            np.save(feature_db_name, feature_db)

def get_latent_vectors(sess, ops, sq_dir):
    train_file_idxs = []
    listDir(sq_dir, train_file_idxs)
    train_file_idxs.sort()
    is_training=False

    #print(len(train_file_idxs))
    batch_num= BATCH_NUM_QUERIES*(1+POSITIVES_PER_QUERY+NEGATIVES_PER_QUERY)
    q_output = []
    for q_index in tqdm(range(len(train_file_idxs)//batch_num)):
        file_names=train_file_idxs[q_index*batch_num:(q_index+1)*(batch_num)]
        queries=load_pc_files(file_names)
        # queries= np.expand_dims(queries,axis=1)
        q1=queries[0:BATCH_NUM_QUERIES]
        q1=np.expand_dims(q1,axis=1)
        #print(q1.shape)

        q2=queries[BATCH_NUM_QUERIES:BATCH_NUM_QUERIES*(POSITIVES_PER_QUERY+1)]
        q2=np.reshape(q2,(BATCH_NUM_QUERIES,POSITIVES_PER_QUERY,NUM_POINTS,3))

        q3=queries[BATCH_NUM_QUERIES*(POSITIVES_PER_QUERY+1):BATCH_NUM_QUERIES*(NEGATIVES_PER_QUERY+POSITIVES_PER_QUERY+1)]
        q3=np.reshape(q3,(BATCH_NUM_QUERIES,NEGATIVES_PER_QUERY,NUM_POINTS,3))
        feed_dict={ops['query']:q1, ops['positives']:q2, ops['negatives']:q3, ops['is_training_pl']:is_training}
        t1 = time.time()
        o1, o2, o3=sess.run([ops['q_vec'], ops['pos_vecs'], ops['neg_vecs']], feed_dict=feed_dict)
        t2 = time.time()
        print("inference time: ", t2-t1)
        o1=np.reshape(o1,(-1,o1.shape[-1]))
        o2=np.reshape(o2,(-1,o2.shape[-1]))
        o3=np.reshape(o3,(-1,o3.shape[-1]))

        out=np.vstack((o1,o2,o3))
        q_output.append(out)

    q_output=np.array(q_output)
    if(len(q_output)!=0):
        q_output=q_output.reshape(-1,q_output.shape[-1])
    #print(q_output.shape)

    #handle edge case
    for q_index in tqdm(range((len(train_file_idxs)//batch_num*batch_num),len(train_file_idxs))):
        file_names=train_file_idxs[q_index]
        queries=load_pc_files([file_names])
        queries= np.expand_dims(queries,axis=1)
        #print(query.shape)
        #exit()
        fake_queries=np.zeros((BATCH_NUM_QUERIES-1,1,NUM_POINTS,3))
        fake_pos=np.zeros((BATCH_NUM_QUERIES,POSITIVES_PER_QUERY,NUM_POINTS,3))
        fake_neg=np.zeros((BATCH_NUM_QUERIES,NEGATIVES_PER_QUERY,NUM_POINTS,3))
        q=np.vstack((queries,fake_queries))
        #print(q.shape)
        feed_dict={ops['query']:q, ops['positives']:fake_pos, ops['negatives']:fake_neg, ops['is_training_pl']:is_training}
        output=sess.run(ops['q_vec'], feed_dict=feed_dict)
        #print(output.shape)
        output=output[0]
        output=np.squeeze(output)
        if (q_output.shape[0]!=0):
            q_output=np.vstack((q_output,output))
        else:
            q_output=output

    #q_output=np.array(q_output)
    #q_output=q_output.reshape(-1,q_output.shape[-1])
    print(q_output.shape)
    print(len(train_file_idxs))
    assert q_output.shape[0] == len(train_file_idxs)
    return q_output

if __name__ == "__main__":
    evaluate()


import numpy as np
import theano
import theano.tensor as T

from .. import init
from .. import nonlinearities

from .base import Layer

from theano.sandbox.cuda.basic_ops import gpu_contiguous
from pylearn2.sandbox.cuda_convnet.filter_acts import FilterActs

__all__ = [
    "CCLayer",
    "Conv2DCCLayer",
    "MaxPool2DCCLayer",
    "ShuffleBC01ToC01BLayer",
    "bc01_to_c01b",
    "ShuffleC01BToBC01Layer",
    "c01b_to_bc01",
    "NINLayer_c01b",
]

if not theano.config.device.startswith("gpu"):
    raise ImportError("requires a GPU to work")

# TODO: make sure to document the limitations and 'best practices'
# (i.e. minibatch size % 128 == 0)
# TODO: see if the 'dimshuffle' logic can be put in the base class instead.

# base class for all layers that use ops from pylearn2.sandbox.cuda_convnet
class CCLayer(Layer):
    pass

class Conv2DCCLayer(CCLayer):
    def __init__(self, incoming, num_filters, filter_size, strides=(1, 1),
                 border_mode=None, untie_biases=False, W=init.Uniform(),
                 b=init.Constant(0.), nonlinearity=nonlinearities.rectify,
                 pad=None, dimshuffle=True, flip_filters=False, partial_sum=1,
                 **kwargs):
        super(Conv2DCCLayer, self).__init__(incoming, **kwargs)
        if nonlinearity is None:
            self.nonlinearity = nonlinearities.identity
        else:
            self.nonlinearity = nonlinearity

        if filter_size[0] != filter_size[1]:
            raise RuntimeError("Conv2DCCLayer only supports square filters, "
                               "but filter_size=(%d, %d)" % filter_size)

        if strides[0] != strides[1]:
            raise RuntimeError("Conv2DCCLayer only supports square strides, "
                               "but strides=(%d, %d)" % strides)

        if num_filters % 16 != 0:
            raise RuntimeError("Conv2DCCLayer requires num_filters to be a "
                               "multiple of 16, but num_filters is "
                               "%d" % num_filters)

        self.num_filters = num_filters
        self.filter_size = filter_size[0]
        self.stride = strides[0]
        self.untie_biases = untie_biases
        self.dimshuffle = dimshuffle
        self.flip_filters = flip_filters
        self.partial_sum = partial_sum

        if border_mode is not None and pad is not None:
            raise RuntimeError("You cannot specify both 'border_mode' and "
                               "'pad'. To avoid ambiguity, please specify "
                               "only one of them.")
        elif border_mode is None and pad is None:
            # no option specified, default to valid mode
            self.pad = 0
        elif border_mode is not None:
            if border_mode == 'valid':
                self.pad = 0
            elif border_mode == 'full':
                self.pad = self.filter_size - 1
            elif border_mode == 'same':
                # only works for odd filter size, but the even filter size case
                # is probably not worth supporting.
                self.pad = (self.filter_size - 1) // 2
            else:
                raise RuntimeError("Unsupported border_mode for "
                                   "Conv2DCCLayer: %s" % border_mode)
        else:
            self.pad = pad

        self.W = self.create_param(W, self.get_W_shape())
        if b is None:
            self.b = None
        elif self.untie_biases:
            output_shape = self.get_output_shape()
            if self.dimshuffle:
                self.b = self.create_param(b, (num_filters, output_shape[2],
                                               output_shape[3]))
            else:
                self.b = self.create_param(b, (num_filters, output_shape[1],
                                               output_shape[2]))
        else:
            self.b = self.create_param(b, (num_filters,))

        self.filter_acts_op = FilterActs(
            stride=self.stride, partial_sum=self.partial_sum, pad=self.pad)

    def get_W_shape(self):
        if self.dimshuffle:
            num_input_channels = self.input_shape[1]
            return (self.num_filters, num_input_channels, self.filter_size,
                    self.filter_size)
        else:
            num_input_channels = self.input_shape[0]
            return (num_input_channels, self.filter_size, self.filter_size,
                    self.num_filters)

    def get_params(self):
        return [self.W] + self.get_bias_params()

    def get_bias_params(self):
        return [self.b] if self.b is not None else []

    def get_output_shape_for(self, input_shape):
        if self.dimshuffle:
            batch_size = input_shape[0]
            input_rows, input_columns = input_shape[2:4]
        else:
            batch_size = input_shape[3]
            input_rows, input_columns = input_shape[1:3]

        output_rows = ((input_rows + 2*self.pad - self.filter_size) //
                       self.stride + 1)
        output_columns = ((input_columns + 2*self.pad - self.filter_size) //
                          self.stride + 1)

        if self.dimshuffle:
            return (batch_size, self.num_filters, output_rows, output_columns)
        else:
            return (self.num_filters, output_rows, output_columns, batch_size)

    def get_output_for(self, input, *args, **kwargs):
        if self.dimshuffle:
            filters = self.W.dimshuffle(1, 2, 3, 0)  # bc01 to c01b
            input = input.dimshuffle(1, 2, 3, 0)  # bc01 to c01b
        else:
            filters = self.W

        if self.flip_filters:
            filters = filters[:, ::-1, ::-1, :]  # flip top-down, left-right

        contiguous_filters = gpu_contiguous(filters)
        contiguous_input = gpu_contiguous(input)
        conved = self.filter_acts_op(contiguous_input, contiguous_filters)

        if self.b is not None:
            if self.untie_biases:
                biases = self.b.dimshuffle(0, 1, 2, 'x')  # c01 to c01b
            else:
                biases = self.b.dimshuffle(0, 'x', 'x', 'x')  # c to c01b
            conved += biases

        conved = self.nonlinearity(conved)

        if self.dimshuffle:
            return conved.dimshuffle(3, 0, 1, 2)  # c01b to bc01
        else:
            return conved

class MaxPool2DCCLayer(CCLayer):
    def __init__(self, incoming, ds, ignore_border=False, strides=None,
                 dimshuffle=True, **kwargs):
        from pylearn2.sandbox.cuda_convnet.pool import MaxPool

        super(MaxPool2DCCLayer, self).__init__(incoming, **kwargs)
        if ds[0] != ds[1]:
            raise RuntimeError("MaxPool2DCCLayer only supports square pooling "
                               "regions, but ds=(%d, %d)" % ds)

        if strides is not None and strides[0] != strides[1]:
            raise RuntimeError("MaxPool2DCCLayer only supports using the same "
                               "stride in both directions, but "
                               "strides=(%d, %d)" % strides)

        # ignore_border argument is for compatibility with MaxPool2DLayer.
        # it is not supported. Borders are never ignored.
        if ignore_border is not False:
            raise RuntimeError("MaxPool2DCCLayer does not support "
                               "ignore_border.")

        self.ds = ds[0]
        if strides is None:
            self.stride = self.ds
        else:
            self.stride = strides[0]
        self.dimshuffle = dimshuffle

        self.pool_op = MaxPool(ds=self.ds, stride=self.stride)

    def get_output_shape_for(self, input_shape):
        if self.dimshuffle:
            batch_size = input_shape[0]
            num_input_channels = input_shape[1]
            input_rows, input_columns = input_shape[2:4]
        else:
            batch_size = input_shape[3]
            num_input_channels = input_shape[0]
            input_rows, input_columns = input_shape[1:3]

        output_rows = int(np.ceil(float(input_rows - self.ds +
                                        self.stride) / self.stride))
        output_columns = int(np.ceil(float(input_columns - self.ds +
                                           self.stride) / self.stride))

        if self.dimshuffle:
            return (batch_size, num_input_channels, output_rows,
                    output_columns)
        else:
            return (num_input_channels, output_rows, output_columns,
                    batch_size)

    def get_output_for(self, input, *args, **kwargs):
        if self.dimshuffle:
            input = input.dimshuffle(1, 2, 3, 0)  # bc01 to c01b

        contiguous_input = gpu_contiguous(input)
        pooled = self.pool_op(contiguous_input)

        if self.dimshuffle:
            return pooled.dimshuffle(3, 0, 1, 2)  # c01b to bc01
        else:
            return pooled

# TODO: crossmapnorm
# from pylearn2.sandbox.cuda_convnet.response_norm import CrossMapNorm

# Helper classes for switching between bc01 and c01b input formats

class ShuffleBC01ToC01BLayer(Layer):
    """
    This layer dimshuffles 4D input for interoperability between c01b and bc01
    ops.
    bc01 (theano) -> c01b (cuda-convnet)
    """
    def get_output_shape_for(self, input_shape):
        return (input_shape[1], input_shape[2], input_shape[3], input_shape[0])

    def get_output_for(self, input, *args, **kwargs):
        return input.dimshuffle(1, 2, 3, 0)

bc01_to_c01b = ShuffleBC01ToC01BLayer  # shortcut

class ShuffleC01BToBC01Layer(Layer):
    """
    This layer dimshuffles 4D input for interoperability between c01b and bc01
    ops.
    c01b (cuda-convnet) -> bc01 (theano)
    """
    def get_output_shape_for(self, input_shape):
        return (input_shape[3], input_shape[0], input_shape[1], input_shape[2])

    def get_output_for(self, input, *args, **kwargs):
        return input.dimshuffle(3, 0, 1, 2)

c01b_to_bc01 = ShuffleC01BToBC01Layer  # shortcut

# c01b versions of other Layer classes

class NINLayer_c01b(Layer):
    """
    This does the same as lasagne.layers.NINLayer, but operates with c01b
    axis arrangement instead of bc01. This reduces the number of shuffles
    and reshapes required and might be faster as a result.
    """
    def __init__(self, incoming, num_units, untie_biases=False,
                 W=init.Uniform(), b=init.Constant(0.),
                 nonlinearity=nonlinearities.rectify, **kwargs):
        super(NINLayer_c01b, self).__init__(incoming, **kwargs)
        if nonlinearity is None:
            self.nonlinearity = nonlinearities.identity
        else:
            self.nonlinearity = nonlinearity

        self.num_units = num_units
        self.untie_biases = untie_biases

        num_input_channels = self.input_shape[0]

        self.W = self.create_param(W, (num_units, num_input_channels))
        if b is None:
            self.b = None
        elif self.untie_biases:
            output_shape = self.get_output_shape()
            self.b = self.create_param(b, (num_units,) + output_shape[1:-1])
        else:
            self.b = self.create_param(b, (num_units,))

    def get_params(self):
        return [self.W] + self.get_bias_params()

    def get_bias_params(self):
        return [self.b] if self.b is not None else []

    def get_output_shape_for(self, input_shape):
        return (self.num_units,) + input_shape[1:]

    def get_output_for(self, input, *args, **kwargs):
        # fc * c01b... = f01b...
        out = T.tensordot(self.W, input, axes=[[1], [0]])

        if self.b is None:
            activation = out
        else:
            if self.untie_biases:
                bias_axes = range(input.ndim - 1) + ['x']
            else:
                bias_axes = [0] + (['x'] * (input.ndim - 1))
            b_shuffled = self.b.dimshuffle(bias_axes)
            activation = out + b_shuffled

        return self.nonlinearity(activation)