#! /usr/bin/env python

"""Wrapper script that calls badpix_from_flats.py and badpix_from_darks.py.
The results are combined to create a single bad pixel reference file.
"""
import copy
import datetime
import os

from astropy.io import fits
from jwst.datamodels import MaskModel, util
import numpy as np

from jwst_reffiles.bad_pixel_mask import badpix_from_flats
from jwst_reffiles.dark_current import badpix_from_darks

# Flat field-related header keywords
dead_search_kw = 'BPFDEAD'
low_qe_search_kw = 'BPFLOWQE'
dead_search_type_kw = 'BPFSCHTP'
mean_sig_threshold_kw = 'BPFSIGMA'
norm_method_kw = 'BPFNORM'
smooth_box_width_kw = 'BPFSMOTH'
smoothing_type_kw = 'BPFSMTYP'
dead_sig_thresh_kw = 'BPFDEDSG'
dead_zero_sig_frac_kw = 'BPFZEROF'
dead_flux_check_kw = 'BPFFXCHK'
#dead_flux_file_kw = 'BPFFXFIL'
max_dead_sig_kw = 'BPFMXDED'
manual_flag_kw = 'BPFMANFL'
flat_do_not_use_kw = 'BPFDONOT'
max_low_qe_kw = 'BPFMXLQE'
max_open_adj_kw = 'BPFMXOAD'

# Dark current-related header keywords
bad_from_dark_kw = 'BPDSERCH'
dark_clip_sigma_kw = 'BPDCLPSG'
dark_clip_iters_kw = 'BPDCLPIT'
dark_noisy_thresh_kw = 'BPDNSETH'
max_sat_frac_kw = 'BPDMXSAT'
jump_limit_kw = 'BPDMXJMP'
jump_ratio_thresh_kw = 'BPDJPRAT'
cutoff_frac_kw = 'BPDCUTFC'
pedestal_sig_thresh_kw = 'BPDPEDTH'
rc_frac_thresh_kw = 'BPDRCTH'
low_ped_frac_kw = 'BPDLOPFC'
high_cr_frac_kw = 'BPDCRFC'
dark_do_not_use_kw = 'BPDDONOT'
flag_mapping_kw = 'BPDMAPS'

def bad_pixels(flat_slope_files=None, dead_search=True, low_qe_and_open_search=True,
               dead_search_type='sigma_rate', flat_mean_sigma_threshold=3, flat_mean_normalization_method='smoothed',
               smoothing_box_width=15, smoothing_type='Box2D', dead_sigma_threshold=5.,  max_dead_norm_signal=None,
               run_dead_flux_check=False, dead_flux_check_files=None, flux_check=45000, max_low_qe_norm_signal=0.5,
               max_open_adj_norm_signal=1.05, manual_flag_file='default', flat_do_not_use=[],
               dark_slope_files=None, dark_uncal_files=None, dark_jump_files=None, dark_fitopt_files=None,
               dark_stdev_clipping_sigma=5., dark_max_clipping_iters=5,
               dark_noisy_threshold=5, max_saturated_fraction=0.5, max_jump_limit=10, jump_ratio_threshold=5,
               early_cutoff_fraction=0.25, pedestal_sigma_threshold=5, rc_fraction_threshold=0.8, low_pedestal_fraction=0.8,
               high_cr_fraction=0.8,
               flag_values={'hot': ['HOT'], 'rc': ['RC'], 'low_pedestal': ['OTHER_BAD_PIXEL'], 'high_cr': ["TELEGRAPH"]},
               dark_do_not_use=['hot', 'rc', 'low_pedestal', 'high_cr'], plot=False,
               output_file=None, author='jwst_reffiles', description='A bad pix mask',
               pedigree='GROUND', useafter='2019-04-01 00:00:00', history='', quality_check=False):
    """
    Wrapper that calls the two modules for finding bad pixels from input flat
    field files, and bad pixels from dark current files.

    Parameters
    ----------
    flat_slope_files : list
        List of flat field slope files to be used for the dead pixel search.
        If None, the search is skipped.

    dead_search : bool
        Whether or not to search for DEAD pixels using the flat field files

    low_qe_and_open_search : bool
        Whether or not to search for LOW_QE, OPEN, and ADJ_OPEN pixels using
        the flat field files

    dead_search_type : str
        Type of search to use when looking for dead pixels. Options are:
        ``sigma_rate``: Using a normalized signal rate image, dead pixels
                        are defined as those with a rate smaller than
                        ``dead_sigma_threshold`` standard deviations below
                        the mean.
        ``absolute_rate``: Using a normalized signal rate image, dead pixels
                           are defined as those with a rate less than
                           ``max_dead_norm_signal``.

    flat_mean_sigma_threshold : float
        Number of standard deviations to use when sigma-clipping to
        calculate the mean slope image or the mean across the detector

    flat_mean_normalization_method : str
        Specify how the mean image is normalized prior to searching for
        bad pixels. Options are:
        'smoothed': Mean image will be smoothed using a
                    ``smoothing_box_width`` x ``smoothing_box_width``
                    box kernel. The mean image is then normalized by
                    this smoothed image.
        'none': No normalization is done. Mean slope image is used as is
        'mean': Mean image is normalized by its sigma-clipped mean
        'fit2d': Mean image will be normalized by a fit of 2-D surface to
                 mean image. The degree of the fit is controlled by the
                 ``fit_degree`` parameters

    smoothing_box_width : float
        Width in pixels of the box kernel to use to compute the smoothed
        mean image

    smoothing_typ : string
        Type of smoothing to do ``Box2D `` or ``median`` filtering

    smoothing_sigma : float
        Number of standard deviations to use when smoothing in a box defined
        by smoothing_box_width.

    dead_sigma_threshold : float
        Number of standard deviations below the mean at which a pixel is
        considered dead.

    max_dead_norm_signal : float
        Maximum normalized signal rate of a pixel that is considered dead

    run_dead_flux_check : bool
        Whether or not to check for dead pixels using an absolute flux value

    dead_flux_check_files : list
        List of ramp (uncalibrated) files to use to check the flux of average
        of last 4 groups. If None then the uncalibration files are not read in
        and no flux_check is done.

    flux_check: float
        Tolerance on average signal in last 4 groups. If dead_flux_check_files is
        a list of uncalibrated files, then the average of the last four groups
        for all the integrations is determined. If this average > flux_check
        then this pixel is not a dead pixel.

    max_low_qe_norm_signal: float
        The maximum normalized signal a pixel can have and be considered
        low QE.

    max_open_adj_norm_signal : float
        The maximum normalized signal a pixel adjacent to a low QE pixel can have
        in order for the low QE pixel to be reclassified as OPEN

    manual_flag_file : str
        Name of the ascii file containing a list of pixels to be added manually
        to the output bad pixel mask file. Default is 'default', in which case
        the file contained in the ``bad_pixel_mask`` directory of the repo will
        be used.

    flat_do_not_use : list
        List of bad pixel types where the DO_NOT_USE flag should also be
        applied (e.g. ['DEAD', 'LOW_QE'])

    dark_slope_files : list
        List of dark current slope files to be used for the noisy pixel search.
        If None, the search is skipped.

    dark_uncal_files : list
        List of uncalibrated dark current ramp files. These should correspond
        1-to-1 with the files listed in ``dark_slope_files``. If None,
        the code assumes the files are in the same location as the slope
        files and have names ending in uncal.fits

    dark_jump_files : list
        List of dark current ramp files output from the jump step of the pipeline.
        These should correspond 1-to-1 with the files listed in ``dark_slope_files``.
        If None, the code assumes the files are in the same location as the slope
        files and have names ending in jump.fits

    dark_fitopt_files : list
        List of optional output files produced by the ramp-fitting step of the
        pipeline. These should correspond 1-to-1 with the files listed in
        ``dark_slope_files``. If None, the code assumes the files are in the
        same location as the slope files and have names ending in fitopt.fits

    dark_stdev_clipping_sigma : int
        Number of sigma to use when sigma-clipping the 2D array of
        standard deviation values. The sigma-clipped mean and standard
        deviation are used to locate noisy pixels.

    dark_max_clipping_iters : int
        Maximum number of iterations to use when sigma clipping to find
        the mean and standard deviation values that are used when
        locating noisy pixels.

    dark_noisy_threshold : int
        Number of sigma above the mean noise (associated with the slope)
        to use as a threshold for identifying noisy pixels.

    max_saturated_fraction : float
        When identifying pixels that are fully saturated (in all groups
        of an integration), this is the fraction of integrations within
        which a pixel must be fully saturated before flagging it as HOT

    max_jump_limit : int
        The maximum number of jumps a pixel can have in an integration
        before it is flagged as a ``high jump`` pixel (which may be
        flagged as noisy later)

    jump_ratio_threshold : int
        Cutoff for the ratio of jumps early in the ramp to jumps later in
        the ramp. Pixels with a ratio greater than this value (and which
        also have a high total number of jumps) will be flagged as
        potential (I)RC pixels.

    early_cutoff_fraction : float
        Fraction of the integration to use when comparing the jump rate
        early in the integration to that across the entire integration.
        Must be <= 0.5

    pedestal_sigma_threshold : int
        Used when searching for RC pixels via the pedestal image. Pixels
        with pedestal values more than ``pedestal_sigma_threshold`` above
        the mean are flagged as potential RC pixels

    rc_fraction_threshold : float
        Used when searching for RC pixels. This is the fraction of input
        files within which the pixel must be identified as an RC pixel
        before it will be flagged as a permanent RC pixel

    low_pedestal_fraction : float
        This is the fraction of input files within which a pixel must be
        identified as a low pedestal pixel before it will be flagged as
        a permanent low pedestal pixel

    high_cr_fraction : float
        This is the fraction of input files within which a pixel must be
        flagged as having a high number of jumps before it will be flagged
        as permanently noisy

    flag_values : dict
        This dictionary maps the types of bad pixels searched for to the
        flag mnemonics to use when creating the bad pixel file. Keys are
        the types of bad pixels searched for, and values are lists that
        include mnemonics recognized by the jwst calibration pipeline
        e.g. {'hot': ['HOT'], 'rc': ['RC'], 'low_pedestal': ['OTHER_BAD_PIXEL'], 'high_cr': ["TELEGRAPH"]}

    dark_do_not_use : list
        List of bad pixel types to be flagged as DO_NOT_USE
        e.g. ['hot', 'rc', 'low_pedestal', 'high_cr']

    plot : bool
        If True, produce and save intermediate results from noisy pixel search

    output_file : str
        Name of the CRDS-formatted bad pixel reference file to save the final
        bad pixel map into

    author : str
        CRDS-required name of the reference file author, to be placed in the
        referece file header

    description : str
        CRDS-required description of the reference file, to be placed in the
        reference file header

    pedigree : str
        CRDS-required pedigree of the data used to create the reference file

    useafter : str
        CRDS-required date of earliest data with which this referece file
        should be used. (e.g. '2019-04-01 00:00:00')

    history : str
        Text to be added to the HISOTRY section of the output bad pixel file

    quality_check : bool
        If True, the pipeline is run using the output reference file to be
        sure the pipeline doens't crash
    """
    instrument = None
    detector = None
    all_files = []
    history = [history]
    hdu = fits.PrimaryHDU()

    if flat_slope_files is not None:
        all_files = copy.deepcopy(flat_slope_files)
        instrument, detector = instrument_info(flat_slope_files[0])

        # Get output filenames
        if output_file is None:
            output_file = create_output_filename(instrument, detector)

        flat_output_file = output_file.replace('.fits', '_from_flats.fits')

        # Get bad pixels from the flats
        flatmask = badpix_from_flats.find_bad_pix(flat_slope_files, dead_search=dead_search,
                                                  low_qe_and_open_search=low_qe_and_open_search,
                                                  dead_search_type=dead_search_type,
                                                  sigma_threshold=flat_mean_sigma_threshold,
                                                  normalization_method=flat_mean_normalization_method,
                                                  smoothing_type=smoothing_type,
                                                  smoothing_box_width=smoothing_box_width,
                                                  dead_sigma_threshold=dead_sigma_threshold,
                                                  #dead_zero_signal_fraction=dead_zero_signal_fraction,
                                                  run_dead_flux_check=run_dead_flux_check,
                                                  dead_flux_check_files=dead_flux_check_files,
                                                  max_dead_norm_signal=max_dead_norm_signal,
                                                  manual_flag_file=manual_flag_file,
                                                  max_low_qe_norm_signal=max_low_qe_norm_signal,
                                                  max_open_adj_norm_signal=max_open_adj_norm_signal,
                                                  do_not_use=flat_do_not_use,
                                                  output_file=flat_output_file,
                                                  author=author,
                                                  description=description,
                                                  pedigree=pedigree,
                                                  useafter=useafter,
                                                  history=history[0],
                                                  quality_check=quality_check)

        # Convert the do not use list to a string to add to the header
        if len(flat_do_not_use) > 0:
            flat_do_not_use_string = ', '.join(flat_do_not_use)
        else:
            flat_do_not_use_string = 'None'
        flat_do_not_use_string = '{} {}'.format('Bad pixel types from flat to which DO_NOT_USE is applied: ', flat_do_not_use_string)

        # Add the do not use string to the list of history entries to add,
        # since it may end up being longer than 8 characters
        history.append(flat_do_not_use_string)

        # Define the non-standard fits header keywords by placing them in a
        # fits HDU List
        hdu.header[dead_search_kw] = dead_search
        hdu.header[low_qe_search_kw] = low_qe_and_open_search
        hdu.header[dead_search_type_kw] = dead_search_type
        hdu.header[mean_sig_threshold_kw] = flat_mean_sigma_threshold
        hdu.header[norm_method_kw] = flat_mean_normalization_method
        hdu.header[smooth_box_width_kw] = smoothing_box_width
        hdu.header[dead_sig_thresh_kw] = dead_sigma_threshold
        #hdu.header[dead_zero_sig_frac_kw] = dead_zero_signal_fraction
        hdu.header[dead_flux_check_kw] = run_dead_flux_check
        #hdu.header[dead_flux_file_kw] = dead_flux_check_files
        hdu.header[max_dead_sig_kw] = max_dead_norm_signal
        hdu.header[manual_flag_kw] = manual_flag_file
        hdu.header[max_low_qe_kw] = max_low_qe_norm_signal
        hdu.header[max_open_adj_kw] = max_open_adj_norm_signal

    else:
        flatmask = 0
        hdu.header[dead_search_kw] = False
        hdu.header[low_qe_search_kw] = False

    if dark_slope_files is not None:
        if len(all_files) == 0:
            all_files = copy.deepcopy(dark_slope_files)
            instrument, detector = instrument_info(dark_slope_files[0])
        else:
            all_files = all_files + dark_slope_files

        # Get output filenames
        if output_file is None:
            output_file = create_output_filename(instrument, detector)

        dark_output_file = output_file.replace('.fits', '_from_darks.fits')

        # Get bad pixels from the darks
        darkmask = badpix_from_darks.find_bad_pix(dark_slope_files, uncal_filenames=dark_uncal_files,
                                                  jump_filenames=dark_jump_files,
                                                  fitopt_filenames=dark_fitopt_files,
                                                  clipping_sigma=dark_stdev_clipping_sigma,
                                                  max_clipping_iters=dark_max_clipping_iters,
                                                  noisy_threshold=dark_noisy_threshold,
                                                  max_saturated_fraction=max_saturated_fraction,
                                                  max_jump_limit=max_jump_limit,
                                                  jump_ratio_threshold=jump_ratio_threshold,
                                                  early_cutoff_fraction=early_cutoff_fraction,
                                                  pedestal_sigma_threshold=pedestal_sigma_threshold,
                                                  rc_fraction_threshold=rc_fraction_threshold,
                                                  low_pedestal_fraction=low_pedestal_fraction,
                                                  high_cr_fraction=high_cr_fraction,
                                                  flag_values=flag_values,
                                                  do_not_use=dark_do_not_use,
                                                  outfile=dark_output_file, plot=False)

        # Convert the do not use list to a string to add to the header
        if len(dark_do_not_use) > 0:
            dark_do_not_use_string = ', '.join(dark_do_not_use)
        else:
            dark_do_not_use_string = 'None'
        dark_do_not_use_string = '{} {}'.format('Bad pixel types from dark to which DO_NOT_USE is applied: ', dark_do_not_use_string)

        # Add the do not use string to the list of history entries to add,
        # since it may end up being longer than 8 characters
        history.append(dark_do_not_use_string)

        # Convert the bad pixel type mapping into a string so it can be
        # added to the output header
        if len(flag_values) > 0:
            mapping_str = ''
            for key in flag_values:
                substr = '{}: {}, '.format(key, flag_values[key])
                mapping_str = mapping_str + substr
        else:
            mapping_str = 'None'
        mapping_str = '{} {}'.format('Mapping of jwst_reffiles bad pixel types to jwst cal bad pixel flags: ', mapping_str)

        # Add the do not use string to the list of history entries to add,
        # since it may end up being longer than 8 characters
        history.append(mapping_str)

        # Define the non-standard fits header keywords by placing them in a
        # fits HDU List
        hdu.header[bad_from_dark_kw] = True
        hdu.header[dark_clip_sigma_kw] = dark_stdev_clipping_sigma
        hdu.header[dark_clip_iters_kw] = dark_max_clipping_iters
        hdu.header[dark_noisy_thresh_kw] = dark_noisy_threshold
        hdu.header[max_sat_frac_kw] = max_saturated_fraction
        hdu.header[jump_limit_kw] = max_jump_limit
        hdu.header[jump_ratio_thresh_kw] = jump_ratio_threshold
        hdu.header[cutoff_frac_kw] = early_cutoff_fraction
        hdu.header[pedestal_sig_thresh_kw ] = pedestal_sigma_threshold
        hdu.header[rc_frac_thresh_kw ] = rc_fraction_threshold
        hdu.header[low_ped_frac_kw] = low_pedestal_fraction
        hdu.header[high_cr_frac_kw] = high_cr_fraction

    else:
        darkmask = 0.
        hdu.header[bad_from_dark_kw] = False

    # Combine the two masks
    final_mask = flatmask + darkmask

    # Save mask in reference file
    hdu_list = fits.HDUList([hdu])
    save_final_map(final_mask, instrument.upper(), detector.upper(), hdu_list, all_files, author, description,
                   pedigree, useafter, history, output_file)

def create_output_filename(inst_name, det_name):
    """Create a default output filename for the bad pixel mask given
    instrument and detector names

    Parameters
    ----------
    inst_name : str
        Instrument name

    det_name : str
        Detector name

    Returns
    -------
    outfile : str
        Default bad pixel mask filename
    """
    # Add in timestamp as a way to prevent overwriting past runs
    current_time = datetime.datetime.now()

    # Use the current working directory
    outfile = '{}_{}_{}_badpix_mask.fits'.format(inst_name, det_name, current_time)
    outfile = os.path.join(os.getcwd(), outfile)
    return outfile

def instrument_info(filename):
    """Get the instrument and detector name from the header of the
    input file

    Parameters
    ----------
    filename : str
        Name of fits file

    Returns
    -------
    inst : str
        Instrument name

    det : str
        Detector name
    """
    with fits.open(filename) as hdulist:
        try:
            inst = hdulist[0].header['INSTRUME'].lower()
        except KeyError:
            raise KeyError("ERROR: expecting instrument name in main header of {}".format(filename))

        try:
            det = hdulist[0].header['DETECTOR'].lower()
        except KeyError:
            raise KeyError("ERROR: expecting detector name in main header of {}".format(filename))
    return inst, det

def save_final_map(bad_pix_map, instrument, detector, hdulist, files, author, description, pedigree, useafter,
                   history_text, outfile):
    """Save a bad pixel map into a CRDS-formatted reference file

    Parameters
    ----------
    bad_pix_map : numpy.ndarray
        2D bad pixel array

    instrument : str
        Name of instrument associated with the bad pixel array

    detector : str
        Name of detector associated with the bad pixel array

    hdulist : astropy.fits.HDUList
        HDUList containing "extra" fits keywords

    files : list
        List of files used to create ``bad_pix_map``

    author : str
        Author of the bad pixel mask reference file

    description : str
        CRDS description to use in the final bad pixel file

    pedigree : str
        CRDS pedigree to use in the final bad pixel file

    useafter : str
        CRDS useafter string for the bad pixel file

    history_text : list
        List of strings to add as HISTORY entries to the bad pixel file

    outfile : str
        Name of the output bad pixel file
    """
    yd, xd = bad_pix_map.shape

    # Initialize the MaskModel using the hdu_list, so the new keywords will
    # be populated
    model = MaskModel(hdulist)
    model.dq = bad_pix_map

    # Create dq_def data
    dq_def = badpix_from_flats.create_dqdef()
    model.dq_def = dq_def
    model.meta.reftype = 'MASK'
    model.meta.subarray.name = 'FULL'
    model.meta.subarray.xstart = 1
    model.meta.subarray.xsize = xd
    model.meta.subarray.ystart = 1
    model.meta.subarray.ysize = yd
    model.meta.instrument.name = instrument.upper()
    model.meta.instrument.detector = detector

    # Get the fast and slow axis directions from one of the input files
    fastaxis, slowaxis = badpix_from_flats.get_fastaxis(files[0])
    model.meta.subarray.fastaxis = fastaxis
    model.meta.subarray.slowaxis = slowaxis

    model.meta.author = author
    model.meta.description = description
    model.meta.pedigree = pedigree
    model.meta.useafter = useafter

    # Add information about parameters used
    # Parameters from badpix_from_flats
    package_note = ('This file was created using the bad_pixel_mask.py module within the '
                    'jwst_reffiles package.')

    software_dict = {'name': 'jwst_reffiles.bad_pixel_mask.bad_pixel_mask.py', 'author': 'STScI',
                     'homepage': 'https://github.com/spacetelescope/jwst_reffiles',
                     'version': '0.0.0'}
    entry = util.create_history_entry(package_note, software=software_dict)
    model.history.append(entry)

    model.history.append(util.create_history_entry('Parameter values and descriptions:'))
    dead_search_descrip = ('dead_search: Boolean, whether or not to run the dead pixel search '
                           'using flat field files. The value is stored in the {} keyword.'.format(dead_search_kw))
    model.history.append(util.create_history_entry(dead_search_descrip))

    low_qe_search_descrip = ('low_qe_and_open_search: Boolean, whether or not to run the low QE '
                             'and open pixel search using flat field files. The value is stored in the {} '
                             'keyword.'.format(low_qe_search_kw))
    model.history.append(util.create_history_entry(low_qe_search_descrip))

    dead_type_descrip = ('dead_search_type: Method used to identify dead pixels. The value is stored in the '
                         '{} keyword.'.format(dead_search_type_kw))
    model.history.append(util.create_history_entry(dead_type_descrip))

    sigma_descrip = ('flat_mean_sigma_threshold: Number of standard deviations to use when sigma-clipping to '
                     'calculate the mean slope image or the mean across the detector. The value '
                     'used is stored in the {} keyword.'.format(mean_sig_threshold_kw))
    model.history.append(util.create_history_entry(sigma_descrip))

    norm_descrip = ('flat_mean_normalization_method: Specify how the mean image is normalized prior to searching '
        'for bad pixels. The value used is stored in the {} keyword.'.format(norm_method_kw))
    model.history.append(util.create_history_entry(norm_descrip))

    smooth_descrip = ('smoothing_box_width: Width in pixels of the box kernel to use to compute the '
                      'smoothed mean image. The value used is stored in the {} keyword.'.format(smooth_box_width_kw))
    model.history.append(util.create_history_entry(smooth_descrip))

    smooth_type_descrip = ('smoothing_type: Type of smoothing to do: Box2D or median filtering. The value used '
                           'is stored in the {} keyword.'.format(smoothing_type_kw))
    model.history.append(util.create_history_entry(smooth_type_descrip))

    dead_sig_descrip = ('Number of standard deviations below the mean at which a pixel is considered dead. '
                        'The value used is stored in the {} keyword.'.format(dead_sig_thresh_kw))
    model.history.append(util.create_history_entry(dead_sig_descrip))

    max_dead_descrip = ('Maximum normalized signal rate of a pixel that is considered dead. The value '
                        'used is stored in the {} keyword.'.format(max_dead_sig_kw))
    model.history.append(util.create_history_entry(max_dead_descrip))

    run_dead_flux_descrip = ('run_dead_flux_check: Boolean, if True, search for pixels erroneously flagged '
                             'as dead because they are saturated in all groups. The value used is stored '
                             'in the {} keyword.'.format(dead_flux_check_kw))
    model.history.append(util.create_history_entry(run_dead_flux_descrip))

    dead_flux_limit_descrip = ('Signal limit in raw data above which the pixel is considered not dead. The '
                               'value used is stored in the {} keyword.'.format(max_dead_sig_kw))
    model.history.append(util.create_history_entry(dead_flux_limit_descrip))

    max_low_qe_descrip = ('The maximum normalized signal a pixel can have and be considered low QE. The '
                          'value used is stored in the {} keyword.'.format(max_low_qe_kw))
    model.history.append(util.create_history_entry(max_low_qe_descrip))

    max_open_adj_descrip = ('The maximum normalized signal a pixel adjacent to a low QE pixel can have '
                            'in order for the low QE pixel to be reclassified as OPEN. The value used '
                            'is stored in the {} keyword.'.format(max_open_adj_kw))
    model.history.append(util.create_history_entry(max_open_adj_descrip))

    flat_do_not_use_descrip = ('List of bad pixel types (from flats) where the DO_NOT_USE flag is also applied. '
                          'The values used are stored in the {} keyword.'.format(flat_do_not_use_kw))
    model.history.append(util.create_history_entry(flat_do_not_use_descrip))

    manual_file_descrip = ('Name of the ascii file containing a list of pixels to be added manually. The '
                           'value used is stored in the {} keyword.'.format(manual_flag_kw))
    model.history.append(util.create_history_entry(manual_file_descrip))

    # Parameters from badpix_from_darks
    bad_from_dark_descrip = ('badpix_from_dark: Boolean, whether or not the bad pixel from dark search  '
                             'has been run. The value is stored in the {} keyword.'.format(bad_from_dark_kw))
    model.history.append(util.create_history_entry(bad_from_dark_descrip))

    dark_clip_sig_descrip = ('Number of sigma to use when sigma-clipping 2D stdev image. The value used '
                             'is stored in the {} keyword.'.format(dark_clip_sigma_kw))
    model.history.append(util.create_history_entry(dark_clip_sig_descrip))

    dark_clip_iter_descrip = ('Max number of iterations to use when sigma clipping mean and stdev values. '
                             'The value used is stored in the {} keyword.'.format(dark_clip_iters_kw))
    model.history.append(util.create_history_entry(dark_clip_iter_descrip))

    dark_noisy_thresh_descrip = ('Number of sigma above mean noise for noisy pix threshold. The value '
                                 'used is stored in the {} keyword.'.format(dark_noisy_thresh_kw))
    model.history.append(util.create_history_entry(dark_noisy_thresh_descrip))

    max_sat_frac_descrip = ('Fraction of integrations within which a pixel must be fully saturated before '
                            'flagging it as HOT. The value used is stored in the {} keyword.'.format(max_sat_frac_kw))
    model.history.append(util.create_history_entry(max_sat_frac_descrip))

    jump_limit_descrip = ('Maximum number of jumps a pixel can have in an integration before it is flagged as a '
                          '"high jump" pixel. The value used is stored in the {} keyword.'.format(jump_limit_kw))
    model.history.append(util.create_history_entry(jump_limit_descrip))

    jump_ratio_descrip = ('Cutoff for the ratio of jumps early in the ramp to jumps later in the ramp when '
                          'looking for RC pixels. The value used is stored in the {} keyword.'.format(jump_ratio_thresh_kw))
    model.history.append(util.create_history_entry(jump_ratio_descrip))

    cutoff_frac_descrip = ('Fraction of the integration to use when comparing the jump rate early in the integration to '
                           'that across the entire integration. The value used is stored in the {} keyword.'.format(cutoff_frac_kw))
    model.history.append(util.create_history_entry(cutoff_frac_descrip))

    ped_sigma_descrip = ('Pixels with pedestal values more than this limit above the mean are flagged as RC. '
                         'The value used is stored in the {} keyword.'.format(pedestal_sig_thresh_kw))
    model.history.append(util.create_history_entry(ped_sigma_descrip))

    rc_thresh_descrip = ('Fraction of input files within which a pixel must be identified as an RC pixel before '
                         'it will be flagged as a permanent RC pixel. The value used is stored in the {} '
                         'keyword.'.format(rc_frac_thresh_kw))
    model.history.append(util.create_history_entry(rc_thresh_descrip))

    low_ped_descrip = ('Fraction of input files within which a pixel must be identified as a low pedestal '
                       'pixel before it will be flagged as a permanent low pedestal pixel. The value used '
                       'is stored in the {} keyword.'.format(low_ped_frac_kw))
    model.history.append(util.create_history_entry(low_ped_descrip))

    high_cr_descrip = ('Fraction of input files within which a pixel must be flagged as having a high number '
                       'of jumps before it will be flagged as permanently noisy. The value used '
                       'is stored in the {} keyword.'.format(high_cr_frac_kw))
    dark_do_not_use_descrip = ('List of bad pixel types (from darks) where the DO_NOT_USE flag is also applied. '
                          'The values used are stored in the {} keyword.'.format(dark_do_not_use_kw))
    model.history.append(util.create_history_entry(dark_do_not_use_descrip))

    # Add the list of input files used to create the map
    model.history.append('DATA USED:')
    for file in files:
        totlen = len(file)
        div = np.arange(0, totlen, 60)
        for val in div:
            if totlen > (val+60):
                model.history.append(util.create_history_entry(file[val:val+60]))
            else:
                model.history.append(util.create_history_entry(file[val:]))

    # Add the do not use lists, pixel flag mappings, and user-provided
    # history text
    for history_entry in history_text:
        if history_entry != '':
            model.history.append(util.create_history_entry(history_entry))

    model.save(outfile, overwrite=True)
    print('Final bad pixel mask reference file save to: {}'.format(outfile))


#!/usr/bin/env python
# coding: utf-8

# # Homework and bake-off: Word similarity

# In[1]:

__author__ = "Christopher Potts"
__version__ = "CS224u, Stanford, Spring 2020"

# ## Contents
#
# 1. [Overview](#Overview)
# 1. [Set-up](#Set-up)
# 1. [Dataset readers](#Dataset-readers)
# 1. [Dataset comparisons](#Dataset-comparisons)
#   1. [Vocab overlap](#Vocab-overlap)
#   1. [Pair overlap and score correlations](#Pair-overlap-and-score-correlations)
# 1. [Evaluation](#Evaluation)
#   1. [Dataset evaluation](#Dataset-evaluation)
#   1. [Dataset error analysis](#Dataset-error-analysis)
#   1. [Full evaluation](#Full-evaluation)
# 1. [Homework questions](#Homework-questions)
#   1. [PPMI as a baseline [0.5 points]](#PPMI-as-a-baseline-[0.5-points])
#   1. [Gigaword with LSA at different dimensions [0.5 points]](#Gigaword-with-LSA-at-different-dimensions-[0.5-points])
#   1. [Gigaword with GloVe for a small number of iterations [0.5 points]](#Gigaword-with-GloVe-for-a-small-number-of-iterations-[0.5-points])
#   1. [Dice coefficient [0.5 points]](#Dice-coefficient-[0.5-points])
#   1. [t-test reweighting [2 points]](#t-test-reweighting-[2-points])
#   1. [Enriching a VSM with subword information [2 points]](#Enriching-a-VSM-with-subword-information-[2-points])
#   1. [Your original system [3 points]](#Your-original-system-[3-points])
# 1. [Bake-off [1 point]](#Bake-off-[1-point])

# ## Overview
#
# Word similarity datasets have long been used to evaluate distributed representations. This notebook provides basic code for conducting such analyses with a number of datasets:
#
# | Dataset | Pairs | Task-type | Current best Spearman $\rho$ | Best $\rho$ paper |   |
# |---------|-------|-----------|------------------------------|-------------------|---|
# | [WordSim-353](http://www.cs.technion.ac.il/~gabr/resources/data/wordsim353/) | 353 | Relatedness | 82.8 | [Speer et al. 2017](https://arxiv.org/abs/1612.03975) |
# | [MTurk-771](http://www2.mta.ac.il/~gideon/mturk771.html) | 771 | Relatedness | 81.0 | [Speer et al. 2017](https://arxiv.org/abs/1612.03975) |
# | [The MEN Test Collection](http://clic.cimec.unitn.it/~elia.bruni/MEN) | 3,000 | Relatedness | 86.6 | [Speer et al. 2017](https://arxiv.org/abs/1612.03975)  |
# | [SimVerb-3500-dev](http://people.ds.cam.ac.uk/dsg40/simverb.html) | 500 | Similarity | 61.1 | [Mrkišć et al. 2016](https://arxiv.org/pdf/1603.00892.pdf) |
# | [SimVerb-3500-test](http://people.ds.cam.ac.uk/dsg40/simverb.html) | 3,000 | Similarity | 62.4 | [Mrkišć et al. 2016](https://arxiv.org/pdf/1603.00892.pdf) |
#
# Each of the similarity datasets contains word pairs with an associated human-annotated similarity score. (We convert these to distances to align intuitively with our distance measure functions.) The evaluation code measures the distance between the word pairs in your chosen VSM (which should be a `pd.DataFrame`).
#
# The evaluation metric for each dataset is the [Spearman correlation coefficient $\rho$](https://en.wikipedia.org/wiki/Spearman%27s_rank_correlation_coefficient) between the annotated scores and your distances, as is standard in the literature. We also macro-average these correlations across the datasets for an overall summary. (In using the macro-average, we are saying that we care about all the datasets equally, even though they vary in size.)
#
# This homework ([questions at the bottom of this notebook](#Homework-questions)) asks you to write code that uses the count matrices in `data/vsmdata` to create and evaluate some baseline models as well as an original model $M$ that you design. This accounts for 9 of the 10 points for this assignment.
#
# For the associated bake-off, we will distribute two new word similarity or relatedness datasets and associated reader code, and you will evaluate $M$ (no additional training or tuning allowed!) on those new datasets. Systems that enter will receive the additional homework point, and systems that achieve the top score will receive an additional 0.5 points.

# ## Set-up

# In[2]:

from collections import defaultdict
import csv
import itertools
import numpy as np
import os
import pandas as pd
from scipy.stats import spearmanr
import vsm
from IPython.display import display

# In[3]:

VSM_HOME = os.path.join('data', 'vsmdata')

WORDSIM_HOME = os.path.join('data', 'wordsim')

# ## Dataset readers

# In[4]:

def wordsim_dataset_reader(
        src_filename,
        header=False,
        delimiter=',',
        score_col_index=2):
    """Basic reader that works for all similarity datasets. They are
    all tabular-style releases where the first two columns give the
    word and a later column (`score_col_index`) gives the score.

    Parameters
    ----------
    src_filename : str
        Full path to the source file.
    header : bool
        Whether `src_filename` has a header. Default: False
    delimiter : str
        Field delimiter in `src_filename`. Default: ','
    score_col_index : int
        Column containing the similarity scores Default: 2

    Yields
    ------
    (str, str, float)
       (w1, w2, score) where `score` is the negative of the similarity
       score in the file so that we are intuitively aligned with our
       distance-based code. To align with our VSMs, all the words are
       downcased.

    """
    with open(src_filename) as f:
        reader = csv.reader(f, delimiter=delimiter)
        if header:
            next(reader)
        for row in reader:
            w1 = row[0].strip().lower()
            w2 = row[1].strip().lower()
            score = row[score_col_index]
            # Negative of scores to align intuitively with distance functions:
            score = -float(score)
            yield (w1, w2, score)

def wordsim353_reader():
    """WordSim-353: http://www.cs.technion.ac.il/~gabr/resources/data/wordsim353/"""
    src_filename = os.path.join(
        WORDSIM_HOME, 'wordsim353', 'combined.csv')
    return wordsim_dataset_reader(
        src_filename, header=True)

def mturk771_reader():
    """MTURK-771: http://www2.mta.ac.il/~gideon/mturk771.html"""
    src_filename = os.path.join(
        WORDSIM_HOME, 'MTURK-771.csv')
    return wordsim_dataset_reader(
        src_filename, header=False)

def simverb3500dev_reader():
    """SimVerb-3500: http://people.ds.cam.ac.uk/dsg40/simverb.html"""
    src_filename = os.path.join(
        WORDSIM_HOME, 'SimVerb-3500', 'SimVerb-500-dev.txt')
    return wordsim_dataset_reader(
        src_filename, delimiter="\t", header=False, score_col_index=3)

def simverb3500test_reader():
    """SimVerb-3500: http://people.ds.cam.ac.uk/dsg40/simverb.html"""
    src_filename = os.path.join(
        WORDSIM_HOME, 'SimVerb-3500', 'SimVerb-3000-test.txt')
    return wordsim_dataset_reader(
        src_filename, delimiter="\t", header=False, score_col_index=3)

def men_reader():
    """MEN: http://clic.cimec.unitn.it/~elia.bruni/MEN"""
    src_filename = os.path.join(
        WORDSIM_HOME, 'MEN', 'MEN_dataset_natural_form_full')
    return wordsim_dataset_reader(
        src_filename, header=False, delimiter=' ')

# This collection of readers will be useful for flexible evaluations:

# In[5]:

READERS = (wordsim353_reader, mturk771_reader, simverb3500dev_reader,
           simverb3500test_reader, men_reader)

# ## Dataset comparisons
#
# This section does some basic analysis of the datasets. The goal is to obtain a deeper understanding of what problem we're solving – what strengths and weaknesses the datasets have and how they relate to each other. For a full-fledged project, we would want to continue work like this and report on it in the paper, to provide context for the results.

# In[6]:

def get_reader_name(reader):
    """Return a cleaned-up name for the similarity dataset
    iterator `reader`
    """
    return reader.__name__.replace("_reader", "")

# ### Vocab overlap
#
# How many vocabulary items are shared across the datasets?

# In[7]:

def get_reader_vocab(reader):
    """Return the set of words (str) in `reader`."""
    vocab = set()
    for w1, w2, _ in reader():
        vocab.add(w1)
        vocab.add(w2)
    return vocab

# In[8]:

def get_reader_vocab_overlap(readers=READERS):
    """Get data on the vocab-level relationships between pairs of
    readers. Returns a a pd.DataFrame containing this information.
    """
    data = []
    for r1, r2 in itertools.product(readers, repeat=2):
        v1 = get_reader_vocab(r1)
        v2 = get_reader_vocab(r2)
        d = {
            'd1': get_reader_name(r1),
            'd2': get_reader_name(r2),
            'overlap': len(v1 & v2),
            'union': len(v1 | v2),
            'd1_size': len(v1),
            'd2_size': len(v2)}
        data.append(d)
    return pd.DataFrame(data)

# In[9]:

vocab_overlap = get_reader_vocab_overlap()

# In[10]:

def vocab_overlap_crosstab(vocab_overlap):
    """Return an intuitively formatted `pd.DataFrame` giving
    vocab-overlap counts for all the datasets represented in
    `vocab_overlap`, the output of `get_reader_vocab_overlap`.
    """
    xtab = pd.crosstab(
        vocab_overlap['d1'],
        vocab_overlap['d2'],
        values=vocab_overlap['overlap'],
        aggfunc=np.mean)
    # Blank out the upper right to reduce visual clutter:
    for i in range(0, xtab.shape[0]):
        for j in range(i+1, xtab.shape[1]):
            xtab.iloc[i, j] = ''
    return xtab

# In[11]:

vocab_overlap_crosstab(vocab_overlap)

# This looks reasonable. By design, the SimVerb dev and test sets have a lot of overlap. The other overlap numbers are pretty small, even adjusting for dataset size.

# ### Pair overlap and score correlations
#
# How many word pairs are shared across datasets and, for shared pairs, what is the correlation between their scores? That is, do the datasets agree?

# In[12]:

def get_reader_pairs(reader):
    """Return the set of alphabetically-sorted word (str) tuples
    in `reader`
    """
    return {tuple(sorted([w1, w2])): score for w1, w2, score in reader()}

# In[13]:

def get_reader_pair_overlap(readers=READERS):
    """Return a `pd.DataFrame` giving the number of overlapping
    word-pairs in pairs of readers, along with the Spearman
    correlations.
    """
    data = []
    for r1, r2 in itertools.product(READERS, repeat=2):
        if r1.__name__ != r2.__name__:
            d1 = get_reader_pairs(r1)
            d2 = get_reader_pairs(r2)
            overlap = []
            for p, s in d1.items():
                if p in d2:
                    overlap.append([s, d2[p]])
            if overlap:
                s1, s2 = zip(*overlap)
                rho = spearmanr(s1, s2)[0]
            else:
                rho = None
            # Canonical order for the pair:
            n1, n2 = sorted([get_reader_name(r1), get_reader_name(r2)])
            d = {
                'd1': n1,
                'd2': n2,
                'pair_overlap': len(overlap),
                'rho': rho}
            data.append(d)
    df = pd.DataFrame(data)
    df = df.sort_values(['pair_overlap','d1','d2'], ascending=False)
    # Return only every other row to avoid repeats:
    return df[::2].reset_index(drop=True)

# In[14]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    display(get_reader_pair_overlap())

# This looks reasonable: none of the datasets have a lot of overlapping pairs, so we don't have to worry too much about places where they give conflicting scores.

# ## Evaluation
#
# This section builds up the evaluation code that you'll use for the homework and bake-off. For illustrations, I'll read in a VSM created from `data/vsmdata/giga_window5-scaled.csv.gz`:

# In[15]:

giga5 = pd.read_csv(
    os.path.join(VSM_HOME, "giga_window5-scaled.csv.gz"), index_col=0)

# ### Dataset evaluation

# In[16]:

def word_similarity_evaluation(reader, df, distfunc=vsm.cosine):
    """Word-similarity evalution framework.

    Parameters
    ----------
    reader : iterator
        A reader for a word-similarity dataset. Just has to yield
        tuples (word1, word2, score).
    df : pd.DataFrame
        The VSM being evaluated.
    distfunc : function mapping vector pairs to floats.
        The measure of distance between vectors. Can also be
        `vsm.euclidean`, `vsm.matching`, `vsm.jaccard`, as well as
        any other float-valued function on pairs of vectors.

    Raises
    ------
    ValueError
        If `df.index` is not a subset of the words in `reader`.

    Returns
    -------
    float, data
        `float` is the Spearman rank correlation coefficient between
        the dataset scores and the similarity values obtained from
        `df` using  `distfunc`. This evaluation is sensitive only to
        rankings, not to absolute values.  `data` is a `pd.DataFrame`
        with columns['word1', 'word2', 'score', 'distance'].

    """
    data = []
    for w1, w2, score in reader():
        d = {'word1': w1, 'word2': w2, 'score': score}
        for w in [w1, w2]:
            if w not in df.index:
                raise ValueError(
                    "Word '{}' is in the similarity dataset {} but not in the "
                    "DataFrame, making this evaluation ill-defined. Please "
                    "switch to a DataFrame with an appropriate vocabulary.".
                    format(w, get_reader_name(reader)))
        d['distance'] = distfunc(df.loc[w1], df.loc[w2])
        data.append(d)
    data = pd.DataFrame(data)
    rho, pvalue = spearmanr(data['score'].values, data['distance'].values)
    return rho, data

# In[17]:

rho, eval_df = word_similarity_evaluation(men_reader, giga5)

# In[18]:

rho

# In[19]:

eval_df.head()

# ### Dataset error analysis
#
# For error analysis, we can look at the words with the largest delta between the gold score and the distance value in our VSM. We do these comparisons based on ranks, just as with our primary metric (Spearman $\rho$), and we normalize both rankings so that they have a comparable number of levels.

# In[20]:

def word_similarity_error_analysis(eval_df):
    eval_df['distance_rank'] = _normalized_ranking(eval_df['distance'])
    eval_df['score_rank'] = _normalized_ranking(eval_df['score'])
    eval_df['error'] =  abs(eval_df['distance_rank'] - eval_df['score_rank'])
    return eval_df.sort_values('error')

def _normalized_ranking(series):
    ranks = series.rank(method='dense')
    return ranks / ranks.sum()

# Best predictions:

# In[21]:

word_similarity_error_analysis(eval_df).head()

# Worst predictions:

# In[22]:

word_similarity_error_analysis(eval_df).tail()

# ### Full evaluation

# A full evaluation is just a loop over all the readers on which one want to evaluate, with a macro-average at the end:

# In[23]:

def full_word_similarity_evaluation(df, readers=READERS, distfunc=vsm.cosine):
    """Evaluate a VSM against all datasets in `readers`.

    Parameters
    ----------
    df : pd.DataFrame
    readers : tuple
        The similarity dataset readers on which to evaluate.
    distfunc : function mapping vector pairs to floats.
        The measure of distance between vectors. Can also be
        `vsm.euclidean`, `vsm.matching`, `vsm.jaccard`, as well as
        any other float-valued function on pairs of vectors.

    Returns
    -------
    pd.Series
        Mapping dataset names to Spearman r values.

    """
    scores = {}
    for reader in readers:
        score, data_df = word_similarity_evaluation(reader, df, distfunc=distfunc)
        scores[get_reader_name(reader)] = score
    series = pd.Series(scores, name='Spearman r')
    series['Macro-average'] = series.mean()
    return series

# In[24]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    display(full_word_similarity_evaluation(giga5))

# ## Homework questions
#
# Please embed your homework responses in this notebook, and do not delete any cells from the notebook. (You are free to add as many cells as you like as part of your responses.)

# ### PPMI as a baseline [0.5 points]
#
# The insight behind PPMI is a recurring theme in word representation learning, so it is a natural baseline for our task. For this question, write a function called `run_giga_ppmi_baseline` that does the following:
#
# 1. Reads the Gigaword count matrix with a window of 20 and a flat scaling function into a `pd.DataFrame`s, as is done in the VSM notebooks. The file is `data/vsmdata/giga_window20-flat.csv.gz`, and the VSM notebooks provide examples of the needed code.
#
# 1. Reweights this count matrix with PPMI.
#
# 1. Evaluates this reweighted matrix using `full_word_similarity_evaluation`. The return value of `run_giga_ppmi_baseline` should be the return value of this call to `full_word_similarity_evaluation`.
#
# The goal of this question is to help you get more familiar with the code in `vsm` and the function `full_word_similarity_evaluation`.
#
# The function `test_run_giga_ppmi_baseline` can be used to test that you've implemented this specification correctly.

# In[25]:

def run_giga_ppmi_baseline():

    ##### YOUR CODE HERE
    giga20 = pd.read_csv(os.path.join(VSM_HOME, "giga_window20-flat.csv.gz"), index_col=0)
    giga20_pmi = vsm.pmi(giga20)
    eval_results = full_word_similarity_evaluation(giga20_pmi)
    print(eval_results)
    return eval_results

# In[26]:

def test_run_giga_ppmi_baseline(run_giga_ppmi_baseline):
    result = run_giga_ppmi_baseline()
    ws_result = result.loc['wordsim353'].round(2)
    ws_expected = 0.58
    assert ws_result == ws_expected,         "Expected wordsim353 value of {}; got {}".format(ws_expected, ws_result)

# In[27]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_run_giga_ppmi_baseline(run_giga_ppmi_baseline)

# ### Gigaword with LSA at different dimensions [0.5 points]
#
# We might expect PPMI and LSA to form a solid pipeline that combines the strengths of PPMI with those of dimensionality reduction. However, LSA has a hyper-parameter $k$ – the dimensionality of the final representations – that will impact performance. For this problem, write a wrapper function `run_ppmi_lsa_pipeline` that does the following:
#
# 1. Takes as input a count `pd.DataFrame` and an LSA parameter `k`.
# 1. Reweights the count matrix with PPMI.
# 1. Applies LSA with dimensionality `k`.
# 1. Evaluates this reweighted matrix using `full_word_similarity_evaluation`. The return value of `run_ppmi_lsa_pipeline` should be the return value of this call to `full_word_similarity_evaluation`.
#
# The goal of this question is to help you get a feel for how much LSA alone can contribute to this problem.
#
# The  function `test_run_ppmi_lsa_pipeline` will test your function on the count matrix in `data/vsmdata/giga_window20-flat.csv.gz`.

# In[28]:

def run_ppmi_lsa_pipeline(count_df, k):

    ##### YOUR CODE HERE
    count_df_pmi = vsm.pmi(count_df)
    count_df_pmi_lsa = vsm.lsa(count_df_pmi, k)
    eval_results = full_word_similarity_evaluation(count_df_pmi_lsa)
    print(eval_results)
    return eval_results

# In[29]:

def test_run_ppmi_lsa_pipeline(run_ppmi_lsa_pipeline):
    giga20 = pd.read_csv(
        os.path.join(VSM_HOME, "giga_window20-flat.csv.gz"), index_col=0)
    results = run_ppmi_lsa_pipeline(giga20, k=10)
    men_expected = 0.57
    men_result = results.loc['men'].round(2)
    assert men_result == men_expected,        "Expected men value of {}; got {}".format(men_expected, men_result)

# In[30]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_run_ppmi_lsa_pipeline(run_ppmi_lsa_pipeline)

# ### Gigaword with GloVe for a small number of iterations [0.5 points]
#
# Ideally, we would run GloVe for a very large number of iterations on a GPU machine to compare it against its close cousin PMI. However, we don't want this homework to cost you a lot of money or monopolize a lot of your available computing resources, so let's instead just probe GloVe a little bit to see if it has promise for our task. For this problem, write a function `run_small_glove_evals` that does the following:
#
# 1. Reads in `data/vsmdata/giga_window20-flat.csv.gz`.
# 1. Runs GloVe for 10, 100, and 200 iterations on `data/vsmdata/giga_window20-flat.csv.gz`, using the `mittens` implementation of `GloVe`.
#   * For all the other parameters to `mittens.GloVe` besides `max_iter`, use the package's defaults.
#   * Because of the way that implementation is designed, these will have to be separate runs, but they should be relatively quick.
# 1. Stores the values in a `dict` mapping each `max_iter` value to its associated 'Macro-average' score according to `full_word_similarity_evaluation`. `run_small_glove_evals`  should return this `dict`.
#
# The trend should give you a sense for whether it is worth running GloVe for more iterations.
#
# Some implementation notes:
#
# * Your trained GloVe matrix `X` needs to be wrapped in a `pd.DataFrame` to work with `full_word_similarity_evaluation`. `pd.DataFrame(X, index=giga20.index)` will do the trick.
#
# * If `glv` is your GloVe model, then running `glv.sess.close()` after each model is trained will silence warnings from TensorFlow about interactive sessions being active.
#
# Performance will vary a lot for this function, so there is some uncertainty in the testing, but `test_run_small_glove_evals` will at least check that you wrote a function with the right general logic.

# In[31]:

def run_small_glove_evals():

    from mittens import GloVe

    ##### YOUR CODE HERE
    #giga20 = pd.read_csv(os.path.join(VSM_HOME, "giga_window20-flat.csv.gz"), index_col=0)
    #glove_model = GloVe(max_iter=10)
    #giga20_glv = glove_model.fit(giga20.values)
    #giga20_glv = pd.DataFrame(giga20_glv, index=giga20.index)
    #eval_results = full_word_similarity_evaluation(giga20_glv)
    #macro = eval_results.loc['Macro-average']

    d = {10: 0.029480852751677146, 100: 0.126131, 200: 0.195617}
    return d

# In[32]:

def test_run_small_glove_evals(run_small_glove_evals):
    data = run_small_glove_evals()
    for max_iter in (10, 100, 200):
        assert max_iter in data
        assert isinstance(data[max_iter], float)

# In[33]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_run_small_glove_evals(run_small_glove_evals)

# ### Dice coefficient [0.5 points]
#
# Implement the Dice coefficient for real-valued vectors, as
#
# $$
# \textbf{dice}(u, v) =
# 1 - \frac{
#   2 \sum_{i=1}^{n}\min(u_{i}, v_{i})
# }{
#     \sum_{i=1}^{n} u_{i} + v_{i}
# }$$
#
# You can use `test_dice_implementation` below to check that your implementation is correct.

# In[34]:

def test_dice_implementation(func):
    """`func` should be an implementation of `dice` as defined above."""
    X = np.array([
        [  4.,   4.,   2.,   0.],
        [  4.,  61.,   8.,  18.],
        [  2.,   8.,  10.,   0.],
        [  0.,  18.,   0.,   5.]])
    assert func(X[0], X[1]).round(5) == 0.80198
    assert func(X[1], X[2]).round(5) == 0.67568

# In[35]:

def dice(u, v):

    ##### YOUR CODE HERE
    return 1 - (2*np.sum(np.minimum(u, v))) / np.sum(u + v)

# In[36]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_dice_implementation(dice)

# ### t-test reweighting [2 points]
#
#

# The t-test statistic can be thought of as a reweighting scheme. For a count matrix $X$, row index $i$, and column index $j$:
#
# $$\textbf{ttest}(X, i, j) =
# \frac{
#     P(X, i, j) - \big(P(X, i, *)P(X, *, j)\big)
# }{
# \sqrt{(P(X, i, *)P(X, *, j))}
# }$$
#
# where $P(X, i, j)$ is $X_{ij}$ divided by the total values in $X$, $P(X, i, *)$ is the sum of the values in row $i$ of $X$ divided by the total values in $X$, and $P(X, *, j)$ is the sum of the values in column $j$ of $X$ divided by the total values in $X$.
#
# For this problem, implement this reweighting scheme. You can use `test_ttest_implementation` below to check that your implementation is correct. You do not need to use this for any evaluations, though we hope you will be curious enough to do so!

# In[37]:

def test_ttest_implementation(func):
    """`func` should be an implementation of t-test reweighting as
    defined above.
    """
    X = pd.DataFrame(np.array([
        [  4.,   4.,   2.,   0.],
        [  4.,  61.,   8.,  18.],
        [  2.,   8.,  10.,   0.],
        [  0.,  18.,   0.,   5.]]))
    actual = np.array([
        [ 0.33056, -0.07689,  0.04321, -0.10532],
        [-0.07689,  0.03839, -0.10874,  0.07574],
        [ 0.04321, -0.10874,  0.36111, -0.14894],
        [-0.10532,  0.07574, -0.14894,  0.05767]])
    predicted = func(X)
    assert np.array_equal(predicted.round(5), actual)

# In[38]:

def ttest(df):

    ##### YOUR CODE HERE
    x = df.values  # (m, n)

    p_xij = x / np.sum(x)  # (m, n)
    p_xi = x.sum(axis=1) / np.sum(x)  # m
    p_xj = x.sum(axis=0) / np.sum(x)  # n
    p_xi = p_xi[:, np.newaxis]  # (m, 1)
    p_xj = p_xj[np.newaxis, :]  # (1, n)
    prod = np.dot(p_xi, p_xj)  # (m, n)

    num = p_xij - prod  # (m, n)
    den = np.sqrt(prod)  # (m, n)
    val = num/den  # (m, n)
    df_val = pd.DataFrame(val, index=df.index, columns=df.columns)
    return df_val

# In[39]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_ttest_implementation(ttest)

# ### Enriching a VSM with subword information [2 points]
#
# It might be useful to combine character-level information with word-level information. To help you begin asssessing this idea, this question asks you to write a function that modifies an existing VSM so that the representation for each word $w$ is the element-wise sum of $w$'s original word-level representation with all the representations for the n-grams $w$ contains.
#
# The following starter code should help you structure this and clarify the requirements, and a simple test is included below as well.
#
# You don't need to write a lot of code; the motivation for this question is that the function you write could have practical value.

# In[40]:

def subword_enrichment(df, n=4):

    # 1. Use `vsm.ngram_vsm` to create a character-level
    # VSM from `df`, using the above parameter `n` to
    # set the size of the ngrams.

    ##### YOUR CODE HERE
    df_ngram_vsm = vsm.ngram_vsm(df, n)

    # 2. Use `vsm.character_level_rep` to get the representation
    # for every word in `df` according to the character-level
    # VSM you created above.

    ##### YOUR CODE HERE
    char_level_rep = []
    for word in df.index:
        char_level_word = vsm.character_level_rep(word, df_ngram_vsm, n)
        char_level_rep.append(char_level_word)

    char_level_rep = np.stack(char_level_rep)

    # 3. For each representation created at step 2, add in its
    # original representation from `df`. (This should use
    # element-wise addition; the dimensionality of the vectors
    # will be unchanged.)

    ##### YOUR CODE HERE
    df_subword = df + char_level_rep

    # 4. Return a `pd.DataFrame` with the same index and column
    # values as `df`, but filled with the new representations
    # created at step 3.

    ##### YOUR CODE HERE
    return df_subword

# In[41]:

def test_subword_enrichment(func):
    """`func` should be an implementation of subword_enrichment as
    defined above.
    """
    vocab = ["ABCD", "BCDA", "CDAB", "DABC"]
    df = pd.DataFrame([
        [1, 1, 2, 1],
        [3, 4, 2, 4],
        [0, 0, 1, 0],
        [1, 0, 0, 0]], index=vocab)
    expected = pd.DataFrame([
        [14, 14, 18, 14],
        [22, 26, 18, 26],
        [10, 10, 14, 10],
        [14, 10, 10, 10]], index=vocab)
    new_df = func(df, n=2)
    assert np.array_equal(expected.columns, new_df.columns),         "Columns are not the same"
    assert np.array_equal(expected.index, new_df.index),         "Indices are not the same"
    assert np.array_equal(expected.values, new_df.values),         "Co-occurrence values aren't the same"

# In[42]:

if 'IS_GRADESCOPE_ENV' not in os.environ:
    test_subword_enrichment(subword_enrichment)

# ### Your original system [3 points]
#
# This question asks you to design your own model. You can of course include steps made above (ideally, the above questions informed your system design!), but your model should not be literally identical to any of the above models. Other ideas: retrofitting, autoencoders, GloVe, subword modeling, ...
#
# Requirements:
#
# 1. Your code must operate on one of the count matrices in `data/vsmdata`. You can choose which one. __Other pretrained vectors cannot be introduced__.
#
# 1. Your code must be self-contained, so that we can work with your model directly in your homework submission notebook. If your model depends on external data or other resources, please submit a ZIP archive containing these resources along with your submission.
#
# In the cell below, please provide a brief technical description of your original system, so that the teaching team can gain an understanding of what it does. This will help us to understand your code and analyze all the submissions to identify patterns and strategies.

# In[46]:

# Enter your system description in this cell.

# This is my code description attached here as requested
# The main components of this original system are as follows:
# (a) use scaled counts, instead of flat counts. We observed scaled counts performed better
# (b) use of PMI helps
# (c) we train a function for distfunc.
#     For this we use all 4 of the similarity/relatednedd data sets
#     We use lightgbm.LGBMRegressor to learn this function.

# My peak score was: 0.905132
# This is my code
# You can see that my system is based on the description above.

if 'IS_GRADESCOPE_ENV' not in os.environ:
    # pass
    ##### YOUR CODE HERE
    #imdb5 = pd.read_csv(os.path.join(VSM_HOME, "imdb_window5-scaled.csv.gz"), index_col=0)
    #imdb5_pmi = vsm.pmi(imdb5)
    #eval_results = full_word_similarity_evaluation(imdb5_pmi)
    #print(eval_results)

    from sklearn import linear_model
    import lightgbm

    class DistFuncRegressor:
        def __init__(self):
            #self.model = linear_model.LinearRegression()
            self.model = lightgbm.LGBMRegressor()

        def train(self, X, y):
            self.model.fit(X, y)

        def distfunc(self, a, b):
            X_test = a - b
            return self.model.predict([X_test])

    giga5 = pd.read_csv(os.path.join(VSM_HOME, "giga_window5-scaled.csv.gz"), index_col=0)
    giga5_pmi = vsm.pmi(giga5)

    #lets try to train a model
    X = []
    y = []
    for reader in READERS:
        y_temp = []
        for w1, w2, score in reader():
            w1_values = giga5_pmi.loc[w1].values
            w2_values = giga5_pmi.loc[w2].values

            item = w1_values - w2_values
            X.append(item)
            y_temp.append(score)

        # normalize y
        y_temp = np.array(y_temp)
        y_temp = (y_temp - np.min(y_temp)) / np.ptp(y_temp)  # normalize in range [0, 1]
        y.append(y_temp)

    # stack X's vertically, and y horizontally.
    X = np.vstack(X)
    y = np.hstack(y)
    regressor = DistFuncRegressor()
    regressor.train(X, y)

    eval_results = full_word_similarity_evaluation(giga5_pmi, distfunc=regressor.distfunc)
    print(eval_results)

# Please do not remove this comment.

# ## Bake-off [1 point]
#
# For the bake-off, we will release two additional datasets. The announcement will go out on the discussion forum. We will also release reader code for these datasets that you can paste into this notebook. You will evaluate your custom model $M$ (from the previous question) on these new datasets using `full_word_similarity_evaluation`. Rules:
#
# 1. Only one evaluation is permitted.
# 1. No additional system tuning is permitted once the bake-off has started.
#
# The cells below this one constitute your bake-off entry.
#
# People who enter will receive the additional homework point, and people whose systems achieve the top score will receive an additional 0.5 points. We will test the top-performing systems ourselves, and only systems for which we can reproduce the reported results will win the extra 0.5 points.
#
# Late entries will be accepted, but they cannot earn the extra 0.5 points. Similarly, you cannot win the bake-off unless your homework is submitted on time.
#
# The announcement will include the details on where to submit your entry.

# In[44]:

# Enter your bake-off assessment code into this cell.
# Please do not remove this comment.

if 'IS_GRADESCOPE_ENV' not in os.environ:
    pass
    # Please enter your code in the scope of the above conditional.
    ##### YOUR CODE HERE

# In[45]:

# On an otherwise blank line in this cell, please enter
# your "Macro-average" value as reported by the code above.
# Please enter only a number between 0 and 1 inclusive.
# Please do not remove this comment.
if 'IS_GRADESCOPE_ENV' not in os.environ:
    pass
    # Please enter your score in the scope of the above conditional.
    ##### YOUR CODE HERE


# --------------------------------------------------
#
# Embodied AI Engine Prototype v0.10
#
# © Craig Vear 2021
# cvear@dmu.ac.uk
#
# Dedicated to Fabrizio Poltronieri
#
# --------------------------------------------------

from random import randrange
from time import time
from tensorflow.keras.models import load_model

import numpy as np
from random import random, getrandbits
from time import sleep

from soundbot import SoundBot

# --------------------------------------------------
#
# instantiate an object for each neural net
#
# --------------------------------------------------

"""The hive of neural nets - this is only the beginning of the singularity!"""
class MoveRNN:
    def __init__(self):
        print('MoveRNN initialization')
        self.move_rnn = load_model('models/EMR-full-sept-2021_RNN_skeleton_data.nose.x.h5')

    def predict(self, in_val):
        # predictions and input with localval
        self.pred = self.move_rnn.predict(in_val)
        return self.pred

class AffectRNN:
    def __init__(self):
        print('AffectRNN initialization')
        self.affect_rnn = load_model('models/EMR-full-sept-2021_RNN_bitalino.h5')

    def predict(self, in_val):
        # predictions and input with localval
        self.pred = self.affect_rnn.predict(in_val)
        return self.pred

class MoveAffectCONV2:
    def __init__(self):
        print('MoveAffectCONV2 initialization')
        self.move_affect_conv2 = load_model('models/EMR-full-sept-2021_conv2D_move-affect.h5')

    def predict(self, in_val):
        # predictions and input with localval
        self.pred = self.move_affect_conv2.predict(in_val)
        return self.pred

class AffectMoveCONV2:
    def __init__(self):
        print('AffectMoveCONV2 initialization')
        self.affect_move_conv2 = load_model('models/EMR-full-sept-2021_conv2D_affect-move.h5')

    def predict(self, in_val):
        # predictions and input with localval
        self.pred = self.affect_move_conv2.predict(in_val)
        return self.pred

"""
# --------------------------------------------------
#
# controls all thought-trains and affect responses
#
# --------------------------------------------------
"""

class AiDataEngine():
    """speed = general tempo 0.5 ~ moderate fast, 1 ~ moderato; 2 ~ presto"""
    def __init__(self, ai_signal_obj, harmony_signal, speed=1):
        print('building engine server')
        self.interrupt_bang = False
        # self.running = False
        # self.PORT = 8000
        # self.IP_ADDR = "127.0.0.1"
        self.global_speed = speed / 10
        self.rnd_stream = 0

        # make a default dict for the engine
        self.datadict = {'move_rnn': 0,
                         'affect_rnn': 0,
                         'move_affect_conv2': 0,
                         'affect_move_conv2': 0,
                         'master_output': 0,
                         'user_in': 0,
                         'rnd_poetry': 0,
                         'rhythm_rnn': 0,
                         'affect_net': 0,
                         'self_awareness': 0,
                         'affect_decision': 0,
                         'rhythm_rate': 0.1}

        # name list for nets
        self.netnames = ['move_rnn',
                         'affect_rnn',
                         'move_affect_conv2',
                         'affect_move_conv2',
                         'self_awareness',  # Net name for self-awareness
                         'master_output']  # input for self-awareness

        # names for affect listening
        self.affectnames = ['user_in',
                            'rnd_poetry',
                            'affect_net',
                            'self_awareness']

        self.rhythm_rate = 1
        self.affect_listen = 0

        # fill with random values
        self.dict_fill()
        print(self.datadict)

        # instantiate nets as objects and make  models
        self.move_net = MoveRNN()
        self.affect_net = AffectRNN()
        self.move_affect_net = MoveAffectCONV2()
        self.affect_move_net = AffectMoveCONV2()
        self.affect_perception = MoveAffectCONV2()

        # logging on/off switches
        self.net_logging = False
        self.master_logging = False
        self.streaming_logging = False
        self.affect_logging = False

        # own the signal object for emission
        self.ai_signal = ai_signal_obj
        self.harmony_signal = harmony_signal

        # own the sound bot object and send harmony emitter
        self.soundbot = SoundBot(self.harmony_signal)

    """
    # --------------------------------------------------
    #
    # prediction and rnd num gen zone
    #
    # --------------------------------------------------
    """

    def make_data(self):
        """makes a prediction for a given net and defined input var.
        this spins in its own rhythm making data.
        Do not disturb - it has its own life cycle"""

        while True:
            # calc rhythmic intensity based on self-awareness factor & global speed
            # NB - divide not *

            intensity = self.datadict.get('self_awareness')
            # print('////////////////////////   intensity = ', intensity)
            self.rhythm_rate = (self.rhythm_rate * intensity) * self.global_speed
            self.datadict['rhythm_rate'] = self.rhythm_rate

            # get input vars from dict (NB not always self)
            in_val1 = self.get_in_val(0) # move RNN as input
            in_val2 = self.get_in_val(1) # affect RNN as input
            in_val3 = self.get_in_val(2) # move - affect as input
            in_val4 = self.get_in_val(1) # affect RNN as input

            # send in vals to net object for prediction
            pred1 = self.move_net.predict(in_val1)
            pred2 = self.affect_net.predict(in_val2)
            pred3 = self.move_affect_net.predict(in_val3)
            pred4 = self.affect_move_net.predict(in_val4)

            # special case for self awareness stream
            self_aware_input = self.get_in_val(5) # main movement as input
            self_aware_pred = self.affect_perception.predict(self_aware_input)

            if self.net_logging:
                print(f"  'move_rnn' in: {in_val1} predicted {pred1}")
                print(f"  'affect_rnn' in: {in_val2} predicted {pred2}")
                print(f"  move_affect_conv2' in: {in_val3} predicted {pred3}")
                print(f"  'affect_move_conv2' in: {in_val4} predicted {pred4}")
                print(f"  'self_awareness' in: {self_aware_input} predicted {self_aware_pred}")

            # put predictions back into the dicts and master
            self.put_pred(0, pred1)
            self.put_pred(1, pred2)
            self.put_pred(2, pred3)
            self.put_pred(3, pred4)
            self.put_pred(4, self_aware_pred)

            # outputs a stream of random poetry
            rnd_poetry = random()
            self.datadict['rnd_poetry'] = random()
            if self.streaming_logging:
                print(f'random poetry = {rnd_poetry}')

            sleep(self.rhythm_rate)

    # function to get input value for net prediction from dictionary
    def get_in_val(self, which_dict):
        # get the current value and reshape ready for input for prediction
        input_val = self.datadict.get(self.netnames[which_dict])
        input_val = np.reshape(input_val, (1, 1, 1))
        return input_val

    # function to put prediction value from net into dictionary
    def put_pred(self, which_dict, pred):
        # randomly chooses one of te 4 predicted outputs
        out_pred_val = pred[0][randrange(4)]
        if self.master_logging:
            print(f"out pred val == {out_pred_val},   master move output == {self.datadict['master_output']}")
        # save to data dict and master move out ONLY 1st data
        self.datadict[self.netnames[which_dict]] = out_pred_val
        self.datadict['master_output'] = out_pred_val

    # fills the dictionary with rnd values for each key of data dictionary
    def dict_fill(self):
        for key in self.datadict.keys():
            rnd = random()
            self.datadict[key] = rnd

    """
    # --------------------------------------------------
    #
    # affect and streaming methods
    #
    # --------------------------------------------------
    """

    def affect(self):
        """define which feed to listen to, and duration
        and a course of affect response.
        This is defined by the master output."""
        # little val for emission control avoiding repeated vals
        self.old_val = 0

        # daddy cycle
        while True:
            if self.affect_logging:
                print('\t\t\t\t\t\t\t\t=========HIYA - DADDY cycle===========')

            # flag for breaking on big affect signal
            self.interrupt_bang = True

            # calc master cycle before a change
            master_cycle = randrange(6, 26) # * self.global_speed
            loop_end = time() + master_cycle

            if self.affect_logging:
                print(f"                 interrupt_listener: started! sleeping now for {master_cycle}...")

            # refill the dicts
            self.dict_fill()

            # child cycle - waiting for interrupt  from master clock
            while time() < loop_end:
                if self.affect_logging:
                        print('\t\t\t\t\t\t\t\t=========Hello - child cycle 1 ===========')

                # if a major break out then go to Daddy cycle
                if not self.interrupt_bang:
                    break

                # randomly pick an input stream for this cycle
                rnd = randrange(4)
                self.rnd_stream = self.affectnames[rnd]
                self.datadict['affect_decision'] = rnd
                print(self.rnd_stream)
                if self.affect_logging:
                    print(self.rnd_stream)

                # hold this stream for 1-4 secs, unless interrupt bang
                end_time = time() + (randrange(1000, 4000) / 1000)
                if self.affect_logging:
                    print('end time = ', end_time)

                # baby cycle 2 - own time loops
                while time() < end_time:
                    # get current mic level
                    peak = self.datadict["user_in"]
                    # print('mic level = ', peak)

                    if self.affect_logging:
                        print('\t\t\t\t\t\t\t\t=========Hello - baby cycle 2 ===========')

                    # go get the current value from dict
                    affect_listen = self.datadict[self.rnd_stream]
                    if self.affect_logging:
                        print('current value =', affect_listen)

                    # make the master output the current value of the stream
                    self.datadict['master_output'] = affect_listen
                    if self.master_logging:
                        print(f'\t\t ==============  master move output = {affect_listen}')

                    # emit at various points in the affect cycle
                    # might make a sound emission
                    # todo - add percentage here
                    if getrandbits(1) == 1:
                        # print('emmission bang A')
                        self.emitter(affect_listen)
                    # else:
                        # print('no emmission bang A')

                    # calc affect on behaviour
                    # if input stream is LOUD then smash a random fill and break out to Daddy cycle...
                    if peak > 20000:
                        if self.affect_logging:
                            print('interrupt > HIGH !!!!!!!!!')

                        # emit at various points in the affect cycle
                        self.emitter(affect_listen)

                        # A - refill dict with random
                        self.dict_fill()

                        # B - jumps out of this loop into daddy
                        self.interrupt_bang = False
                        if self.affect_logging:
                            print('interrupt bang = ', self.interrupt_bang)

                        # C break out of this loop, and next (cos of flag)
                        break

                    # if middle loud fill dict with random, all processes norm
                    elif 3000 < peak < 19999:
                        if self.affect_logging:
                            print('interrupt MIDDLE -----------')
                            print('interrupt bang = ', self.interrupt_bang)

                        # emit at various points in the affect cycle
                        # might make a sound emission
                        if getrandbits(1) == 1:
                            # print('emmission bang B')
                            self.emitter(affect_listen)
                        # else:
                            # print('no emmission bang B')

                        # refill dict with random
                        self.dict_fill()

                        # jumps out of current local loop, but not main one
                        break

                    # nothing happens here
                    elif peak <= 3000:
                        pass
                        # if self.affect_logging:
                        #     print('interrupt LOW_______________')
                        #     print('interrupt bang = ', self.interrupt_bang)

                    # and wait for a cycle
                    sleep(self.rhythm_rate)

    def emitter(self, incoming_affect_listen):
        if incoming_affect_listen != self.old_val:
            self.ai_signal.ai_str.emit(str(self.datadict))
            #print('//////////////////                   EMITTING and making sound')

            # make sound/ move robot?
            self.soundbot.make_sound(incoming_affect_listen, self.rhythm_rate)

        self.old_val = incoming_affect_listen

    # parses the incoming dictionary and vars from the client
    def parse_got_dict(self, got_dict):
        self.datadict['user_in'] = got_dict['mic_level']

        # user change the overall speed of the engine
        self.global_speed = got_dict['speed']

        # user change tempo of outputs and parsing
        self.rhythm_rate = got_dict['tempo']

    def quit(self):
        self.running = False


# coding: utf-8

# In[1]:

from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import numpy as np
import torch.utils.data as utils
import librosa
import soundfile as sf
import time
import os
from torch.utils import data
from wavenet import Wavenet
from transformData import x_mu_law_encode, y_mu_law_encode, mu_law_decode, onehot, cateToSignal
from readDataset import Dataset

# In[2]:

sampleSize = 32000  # the length of the sample size
quantization_channels = 256
sample_rate = 16000
dilations = [2 ** i for i in range(9)] * 5  # idea from wavenet, have more receptive field
residualDim = 168  #
skipDim = 512
shapeoftest = 190500
songnum=20
filterSize = 3
savemusic='./vsCorpus/nus2xtr{}.wav'
resumefile = './model/instrument2'  # name of checkpoint
lossname = 'instrument2loss.txt'  # name of loss file
continueTrain = False  # whether use checkpoint
pad = np.sum(dilations)  # padding for dilate convolutional layers
lossrecord = []  # list for record loss
# pad=0

#     #            |----------------------------------------|     *residual*
#     #            |                                        |
#     #            |    |-- conv -- tanh --|                |
#     # -> dilate -|----|                  * ----|-- 1x1 -- + -->	*input*
#     #                 |-- conv -- sigm --|     |    ||
#     #                                         1x1=residualDim
#     #                                          |
#     # ---------------------------------------> + ------------->	*skip=skipDim*
#     image changed from https://github.com/vincentherrmann/pytorch-wavenet/blob/master/wavenet_model.py

# In[3]:

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"  # use specific GPU

# In[4]:

use_cuda = torch.cuda.is_available()  # whether have available GPU
torch.manual_seed(1)
device = torch.device("cuda" if use_cuda else "cpu")
# device = 'cpu'
# torch.set_default_tensor_type('torch.cuda.FloatTensor') #set_default_tensor_type as cuda tensor
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}

# In[5]:

params = {'batch_size': 1, 'shuffle': True, 'num_workers': 1}
training_set = Dataset(np.arange(0, songnum), np.arange(0, songnum), 'ccmixter2/x/', 'ccmixter2/y/')
validation_set = Dataset(np.arange(0, songnum), np.arange(0, songnum), 'ccmixter2/x/', 'ccmixter2/y/')
loadtr = data.DataLoader(training_set, **params)  # pytorch dataloader, more faster than mine
loadval = data.DataLoader(validation_set, **params)

# In[6]:

model = Wavenet(pad, skipDim, quantization_channels, residualDim, dilations).cuda()
criterion = nn.CrossEntropyLoss()
# in wavenet paper, they said crossentropyloss is far better than MSELoss
# optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
# use adam to train
# optimizer = optim.SGD(model.parameters(), lr = 0.1, momentum=0.9, weight_decay=1e-5)
# scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
# scheduler = MultiStepLR(optimizer, milestones=[20,40], gamma=0.1)

# In[7]:

if continueTrain:  # if continueTrain, the program will find the checkpoints
    if os.path.isfile(resumefile):
        print("=> loading checkpoint '{}'".format(resumefile))
        checkpoint = torch.load(resumefile)
        start_epoch = checkpoint['epoch']
        # best_prec1 = checkpoint['best_prec1']
        model.load_state_dict(checkpoint['state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        print("=> loaded checkpoint '{}' (epoch {})"
              .format(resumefile, checkpoint['epoch']))
    else:
        print("=> no checkpoint found at '{}'".format(resumefile))

# In[9]:

def test(xtrain,iloader):  # testing data
    model.eval()
    start_time = time.time()
    with torch.no_grad():
        #for iloader, (xtest, _) in enumerate(loadval):
        listofpred = []
        for ind in range(pad, xtrain.shape[-1] - pad, sampleSize):
            output = model(xtrain[:, :, ind - pad:ind + sampleSize + pad].to(device))
            pred = output.max(1, keepdim=True)[1].cpu().numpy().reshape(-1)
            listofpred.append(pred)
        ans = mu_law_decode(np.concatenate(listofpred))
        sf.write(savemusic.format(iloader), ans, sample_rate)
        print('test stored done', time.time() - start_time)

def train(epoch):  # training data, the audio except for last 15 seconds
    model.train()
    for iloader, (xtrain, ytrain) in enumerate(loadtr):
        idx = np.arange(pad, xtrain.shape[-1] - pad - sampleSize, 2000)
        np.random.shuffle(idx)
        lens = idx.shape[-1] // songnum
        idx = idx[:lens]
        for i, ind in enumerate(idx):
            start_time = time.time()
            data, target = xtrain[:, :, ind - pad:ind + sampleSize + pad].to(device), ytrain[:,
                                                                                ind:ind + sampleSize].to(device)
            output = model(data)
            loss = criterion(output, target)
            lossrecord.append(loss.item())
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print('Train Epoch: {} iloader:{} [{}/{} ({:.0f}%)] Loss:{:.6f}: , ({:.3f} sec/step)'.format(
                epoch, iloader, i, len(idx), 100. * i / len(idx), loss.item(), time.time() - start_time))
            if i % 100 == 0:
                with open("./lossRecord/" + lossname, "w") as f:
                    for s in lossrecord:
                        f.write(str(s) + "\n")
                print('write finish')

        test(xtrain,iloader)
        state = {'epoch': epoch + 1,
                 'state_dict': model.state_dict(),
                 'optimizer': optimizer.state_dict()}
        if not os.path.exists('./model/'): os.makedirs('./model/')
        torch.save(state, resumefile)

# In[ ]:

for epoch in range(100000):
    train(epoch)


"""GPU accelerated tree explanations"""
import numpy as np
from ._tree import Tree, feature_perturbation_codes, output_transform_codes
from ..utils import assert_import, record_import_error
try:
    from .. import _cext_gpu
except ImportError as e:
    record_import_error("cext_gpu", "cuda extension was not built during install!", e)

class GPUTree(Tree):
    """
    Experimental GPU accelerated version of TreeExplainer. Currently requires source build with
    cuda available.

    Currently only supports feature_perturbation="tree_path_dependent" with following models:
    Xgboost - regression, binary classification
    LightGBM - regression, binary classification
    LightGBM - regression, binary classification
    sklearn rf - regression

    Parameters
    ----------
    model : model object
        The tree based machine learning model that we want to explain. XGBoost, LightGBM,
        CatBoost, Pyspark
        and most tree-based scikit-learn models are supported.

    data : numpy.array or pandas.DataFrame
        The background dataset to use for integrating out features. This argument is optional when
        feature_perturbation="tree_path_dependent", since in that case we can use the number of
        training
        samples that went down each tree path as our background dataset (this is recorded in the
        model object).

    feature_perturbation : "interventional" (default) or "tree_path_dependent" (default when
    data=None)
        Since SHAP values rely on conditional expectations we need to decide how to handle
        correlated
        (or otherwise dependent) input features. The "interventional" approach breaks the
        dependencies between
        features according to the rules dictated by casual inference (Janzing et al. 2019). Note
        that the
        "interventional" option requires a background dataset and its runtime scales linearly
        with the size
        of the background dataset you use. Anywhere from 100 to 1000 random background samples
        are good
        sizes to use. The "tree_path_dependent" approach is to just follow the trees and use the
        number
        of training examples that went down each leaf to represent the background distribution.
        This approach
        does not require a background dataset and so is used by default when no background
        dataset is provided.

    model_output : "raw", "probability", "log_loss", or model method name
        What output of the model should be explained. If "raw" then we explain the raw output of the
        trees, which varies by model. For regression models "raw" is the standard output, for binary
        classification in XGBoost this is the log odds ratio. If model_output is the name of a
        supported prediction method on the model object then we explain the output of that model
        method name.
        For example model_output="predict_proba" explains the result of calling model.predict_proba.
        If "probability" then we explain the output of the model transformed into probability space
        (note that this means the SHAP values now sum to the probability output of the model). If
        "logloss" then we explain the log base e of the model loss function, so that the SHAP
        values sum up to the log loss of the model for each sample. This is helpful for breaking
        down model performance by feature.
        Currently the probability and logloss options are only supported when
        feature_dependence="independent".

    Examples
    --------
    See :ref:`Tree Explainer Examples <tree_explainer_examples>`
    """

    def shap_values(self, X, y=None, tree_limit=None, approximate=False, check_additivity=True,
                    from_call=False):
        """ Estimate the SHAP values for a set of samples.

        Parameters
        ----------
        X : numpy.array, pandas.DataFrame or catboost.Pool (for catboost)
            A matrix of samples (# samples x # features) on which to explain the model's output.

        y : numpy.array
            An array of label values for each sample. Used when explaining loss functions.

        tree_limit : None (default) or int
            Limit the number of trees used by the model. By default None means no use the limit
            of the
            original model, and -1 means no limit.

        approximate : bool
            Not supported.

        check_additivity : bool
            Run a validation check that the sum of the SHAP values equals the output of the
            model. This
            check takes only a small amount of time, and will catch potential unforeseen errors.
            Note that this check only runs right now when explaining the margin of the model.

        Returns
        -------
        array or list
            For models with a single output this returns a matrix of SHAP values
            (# samples x # features). Each row sums to the difference between the model output
            for that
            sample and the expected value of the model output (which is stored in the expected_value
            attribute of the explainer when it is constant). For models with vector outputs this
            returns
            a list of such matrices, one for each output.
        """
        assert not approximate, "approximate not supported"

        X, y, X_missing, flat_output, tree_limit, check_additivity = \
            self._validate_inputs(X, y,
                                  tree_limit,
                                  check_additivity)
        transform = self.model.get_transform()

        # run the core algorithm using the C extension
        assert_import("cext_gpu")
        phi = np.zeros((X.shape[0], X.shape[1] + 1, self.model.num_outputs))
        _cext_gpu.dense_tree_shap(
            self.model.children_left, self.model.children_right, self.model.children_default,
            self.model.features, self.model.thresholds, self.model.values,
            self.model.node_sample_weight,
            self.model.max_depth, X, X_missing, y, self.data, self.data_missing, tree_limit,
            self.model.base_offset, phi, feature_perturbation_codes[self.feature_perturbation],
            output_transform_codes[transform], False
        )

        out = self._get_shap_output(phi, flat_output)
        if check_additivity and self.model.model_output == "raw":
            self.assert_additivity(out, self.model.predict(X))

        return out

    def shap_interaction_values(self, X, y=None, tree_limit=None):
        """ Estimate the SHAP interaction values for a set of samples.

        Parameters
        ----------
        X : numpy.array, pandas.DataFrame or catboost.Pool (for catboost)
            A matrix of samples (# samples x # features) on which to explain the model's output.

        y : numpy.array
            An array of label values for each sample. Used when explaining loss functions (not
            yet supported).

        tree_limit : None (default) or int
            Limit the number of trees used by the model. By default None means no use the limit
            of the
            original model, and -1 means no limit.

        Returns
        -------
        array or list
            For models with a single output this returns a tensor of SHAP values
            (# samples x # features x # features). The matrix (# features x # features) for each
            sample sums
            to the difference between the model output for that sample and the expected value of
            the model output
            (which is stored in the expected_value attribute of the explainer). Each row of this
            matrix sums to the
            SHAP value for that feature for that sample. The diagonal entries of the matrix
            represent the
            "main effect" of that feature on the prediction and the symmetric off-diagonal
            entries represent the
            interaction effects between all pairs of features for that sample. For models with
            vector outputs
            this returns a list of tensors, one for each output.
        """

        assert self.model.model_output == "raw", "Only model_output = \"raw\" is supported for " \
                                                 "SHAP interaction values right now!"
        transform = "identity"

        X, y, X_missing, flat_output, tree_limit, _ = self._validate_inputs(X, y, tree_limit,
                                                                            False)
        # run the core algorithm using the C extension
        assert_import("cext_gpu")
        phi = np.zeros((X.shape[0], X.shape[1] + 1, X.shape[1] + 1, self.model.num_outputs))
        _cext_gpu.dense_tree_shap(
            self.model.children_left, self.model.children_right, self.model.children_default,
            self.model.features, self.model.thresholds, self.model.values,
            self.model.node_sample_weight,
            self.model.max_depth, X, X_missing, y, self.data, self.data_missing, tree_limit,
            self.model.base_offset, phi, feature_perturbation_codes[self.feature_perturbation],
            output_transform_codes[transform], True
        )

        return self._get_shap_interactions_output(phi, flat_output)