import numpy as np
import sqlite3 as sq
import datetime as dt
import subprocess as sp
import glob as gb
import os

import matplotlib.pyplot as plt

from PyFVCOM.grid import vincenty_distance
from PyFVCOM.read import FileReader
from PyFVCOM.plot import Time, Plotter
from PyFVCOM.stats import calculate_coefficient, rmse

SQL_UNIX_EPOCH = dt.datetime(1970, 1, 1, 0, 0, 0)

"""
Validation of model outputs against in situ data stored and extracted from a database.

This also includes the code to build the databases of time series data sets.

"""

class validation_db():
    """ Work with an SQLite database. """

    def __init__(self, db_name):
        if db_name[-3:] != '.db':
            db_name += '.db'
        self.conn = sq.connect(db_name)
        self.create_table_sql = {}
        self.retrieve_data_sql = {}
        self.c = self.conn.cursor()

    def execute_sql(self, sql_str):
        """
        Execute the given SQL statement.

        Parameters
        ----------
        sql_str : str
            SQL statement to execute.

        Returns
        -------
        results : np.ndarray
            Data from the database which matches the SQL statement.
        """

        self.c.execute(sql_str)

        return self.c.fetchall()

    def make_create_table_sql(self, table_name, col_list):
        """
        Create an SQL table if no such table exists.

        Parameters
        ----------
        table_name : str
            Table name to create.
        col_list : list
            List of column names to add to the table.

        """

        create_str = 'CREATE TABLE IF NOT EXISTS ' + table_name + ' ('
        for this_col in col_list:
            create_str += this_col
            create_str += ', '
        create_str = create_str[0:-2]
        create_str += ');'
        self.create_table_sql['create_' + table_name] = create_str

    def insert_into_table(self, table_name, data):
        """
        Insert data into a table.

        Parameters
        ----------
        table_name : str
            Table name into which to insert the given data.
        data : np.ndarray
            Data to insert into the database.

        """
        no_rows = len(data)
        no_cols = len(data[0])
        qs_string = '('
        for this_x in range(no_cols):
            qs_string += '?,'
        qs_string = qs_string[:-1]
        qs_string += ')'

        if no_rows > 1:
            self.c.executemany('insert or ignore into ' + table_name + ' values ' + qs_string, data)
        elif no_rows == 1:
            self.c.execute('insert into ' + table_name + ' values ' + qs_string, data[0])
        self.conn.commit()

    def select_qry(self, table_name, where_str, select_str='*', order_by_str=None, inner_join_str=None, group_by_str=None):
        """
        Extract data from the database which matches the given SQL query.

        Parameters
        ----------
        table_name : str
            Table name to query.
        where_str : str
            Where statement.
        select_str : str, optional
            Optionally give a set of columns to select.
        order_by_str : str, optional
            Optionally give a set of columns by which to order the results.
        inner_join_str : str, optional
            Optionally give an inner join string.
        group_by_str : str, optional
            Optionally give a string by which to group the results.

        Returns
        -------
        results : np.ndarray
            The data which matches the given query.

        """
        qry_string = 'select ' + select_str + ' from ' + table_name
        if inner_join_str:
            qry_string += ' inner join ' + inner_join_str
        if where_str:
            qry_string += ' where ' + where_str
        if order_by_str:
            qry_string += ' order by ' + order_by_str
        if group_by_str:
            qry_string += ' group by ' + group_by_str

        return self.execute_sql(qry_string)

    def close_conn(self):
        """ Close the connection to the database. """
        self.conn.close()

def dt_to_epochsec(time_to_convert):
    """
    Convert a datetime to our SQL database epoch.

    Parameters
    ----------
    time_to_convert : datetime.datetime
        Datetime to convert.

    Returns
    -------
    epoched : int
        Converted datetime (in seconds).

    """

    return (time_to_convert - SQL_UNIX_EPOCH).total_seconds()

def epochsec_to_dt(time_to_convert):
    """

    Parameters
    ----------
    time_to_convert : int
        Seconds in the SQL database epoch.

    Return
    ------
    unepoched : datetime.datetime.
        Converted time.

    """
    return SQL_UNIX_EPOCH + dt.timedelta(seconds=time_to_convert)

def plot_map(fvcom, tide_db_path, threshold=np.inf, legend=False, **kwargs):
    """
    Plot the tide gauges which fall within the model domain (in space and time) defined by the given FileReader object.

    Parameters
    ----------
    fvcom : PyFVCOM.read.FileReader
        FVCOM model data as a FileReader object.
    tide_db_path : str
        Path to the tidal database.
    threshold : float, optional
        Give a threshold distance (in spherical units) beyond which a gauge is considered too far away.
    legend : bool, optional
        Set to True to add a legend to the plot. Defaults to False.

    Any remaining keyword arguments are passed to PyFVCOM.plot.Plotter.

    Returns
    -------
    plot : PyFVCOM.plot.Plotter
        The Plotter object instance for the map

    """

    tide_db = db_tide(tide_db_path)
    gauge_names, gauge_locations = tide_db.get_gauge_locations(long_names=True)

    gauges_in_domain = []
    fvcom_nodes = []
    for gi, gauge in enumerate(gauge_locations):
        river_index = fvcom.closest_node(gauge, threshold=threshold)
        if river_index:
            gauge_id, gauge_dist  = tide_db.get_nearest_gauge_id(*gauge)
            times, data = tide_db.get_tidal_series(gauge_id, np.min(fvcom.time.datetime), np.max(fvcom.time.datetime))
            if not np.any(data):
                continue

            gauges_in_domain.append(gi)
            fvcom_nodes.append(river_index)

    plot = Plotter(fvcom, **kwargs)
    fx, fy = plot.m(fvcom.grid.lon, fvcom.grid.lat)
    plot.plot_field(-fvcom.grid.h)
    plot.axes.plot(fx[fvcom_nodes], fy[fvcom_nodes], 'ro', markersize=3, zorder=202, label='Model')
    # Add the gauge locations.
    rx, ry = plot.m(gauge_locations[:, 0], gauge_locations[:, 1])
    plot.axes.plot(rx, ry, 'wo', label='Gauges')
    for xx, yy, name in zip(rx, ry, gauge_names[gauges_in_domain]):
        plot.axes.text(xx, yy, name, fontsize=10, rotation=45, rotation_mode='anchor', zorder=203)

    if legend:
        plot.axes.legend(numpoints=1, scatterpoints=1, ncol=2, loc='upper center', fontsize=10)

    return plot

def plot_tides(fvcom, db_name, threshold=500, figsize=(10, 10), **kwargs):
    """
    Plot model and tide gauge data.

    Parameters
    ----------
    fvcom : PyFVCOM.read.FileReader
        FVCOM model data as a FileReader object.
    db_name : str
        Database name to interrogate.
    threshold : float, optional
        Give a threshold distance (in spherical units) to exclude gauges too far from a model node.
    figsize : tuple
        Give a figure size (units are inches).

    Remaining keyword arguments are passed to PyFVCOM.plot.Time.

    Returns
    -------
    time : PyFVCOM.plot.Time
        Time series plot object.
    gauge_obs : dict
        Dictionary with the gauge and model data.

    """

    tide_db = db_tide(db_name)

    # Get all the gauges in the database and find the corresponding model nodes.
    gauge_names, gauge_locations = tide_db.get_gauge_locations(long_names=True)

    gauge_obs = {}
    gauges_in_domain = []
    fvcom_nodes = []
    for gi, gauge in enumerate(gauge_locations):
        river_index = fvcom.closest_node(gauge, threshold=threshold)
        if river_index:
            current_gauge = {}
            current_gauge['gauge_id'], current_gauge['gauge_dist'] = tide_db.get_nearest_gauge_id(*gauge)
            current_gauge['times'], current_gauge['data'] = tide_db.get_tidal_series(current_gauge['gauge_id'],
                                                                                np.min(fvcom.time.datetime),
                                                                                np.max(fvcom.time.datetime))
            if not np.any(current_gauge['data']):
                continue

            current_gauge['lon'], current_gauge['lat'] = gauge_locations[gi, :]

            current_gauge['gauge_clean'] = current_gauge['data'][:, 1] == 0
            current_gauge['gauge_obs_clean'] = {'times': np.copy(current_gauge['times'])[current_gauge['gauge_clean']],
                                                'data': np.copy(current_gauge['data'])[current_gauge['gauge_clean'], 0]}
            current_gauge['rescale_zeta'] = fvcom.data.zeta[:, river_index] - np.mean(fvcom.data.zeta[:, river_index])
            current_gauge['rescale_gauge_obs'] = current_gauge['gauge_obs_clean']['data'] - np.mean(current_gauge['gauge_obs_clean']['data'])

            current_gauge['dates_mod'] = np.isin(fvcom.time.datetime, current_gauge['gauge_obs_clean']['times'])
            current_gauge['dates_obs'] = np.isin(current_gauge['gauge_obs_clean']['times'], fvcom.time.datetime)
            # Skip out if we don't have any coincident data (might simply be a sampling issue) within the model
            # period. We should interpolate here.
            if not np.any(current_gauge['dates_mod']) or not np.any(current_gauge['dates_obs']):
                continue

            current_gauge['r'], current_gauge['p'] = calculate_coefficient(current_gauge['rescale_zeta'][current_gauge['dates_mod']], current_gauge['rescale_gauge_obs'][current_gauge['dates_obs']])
            current_gauge['rms'] = rmse(current_gauge['rescale_zeta'][current_gauge['dates_mod']], current_gauge['rescale_gauge_obs'][current_gauge['dates_obs']])
            current_gauge['std'] = np.std(current_gauge['rescale_zeta'][current_gauge['dates_mod']] - current_gauge['rescale_gauge_obs'][current_gauge['dates_obs']])

            gauges_in_domain.append(gi)
            fvcom_nodes.append(river_index)

            name = gauge_names[gi]
            gauge_obs[name] = current_gauge
            del current_gauge

    tide_db.close_conn()  # tidy up after ourselves

    # Now make a figure of all that data.
    if len(gauge_obs) > 5:
        cols = np.ceil(len(gauge_obs) ** (1.0 / 3)).astype(int) + 1
    else:
        cols = 1
    rows = np.ceil(len(gauge_obs) / cols).astype(int)
    fig = plt.figure(figsize=figsize)
    for count, site in enumerate(sorted(gauge_obs)):
        ax = fig.add_subplot(rows, cols, count + 1)
        time = Time(fvcom, figure=fig, axes=ax, hold=True, **kwargs)
        time.plot_line(gauge_obs[site]['rescale_zeta'], label='Model', color='k')
        # We have to use the raw plot function for the gauge data as the plot_line function assumes we're using model
        # data.
        time.axes.plot(gauge_obs[site]['gauge_obs_clean']['times'], gauge_obs[site]['rescale_gauge_obs'], label='Gauge', color='m')
        # Should add the times of the flagged data here.
        time.axes.set_xlim(fvcom.time.datetime.min(), fvcom.time.datetime.max())
        time.axes.set_ylim(np.min((gauge_obs[site]['rescale_gauge_obs'].min(), gauge_obs[site]['rescale_zeta'].min())),
                           np.max((gauge_obs[site]['rescale_gauge_obs'].max(), gauge_obs[site]['rescale_zeta'].max())))
        time.axes.set_title(site)

    return time, gauge_obs

def _make_normal_tide_series(h_series):
    height_series = h_series - np.mean(h_series)
    return height_series

class db_tide(validation_db):
    """ Create a time series database and query it. """

    def make_bodc_tables(self):
        """ Make the complete set of empty tables for data to be inserted into (as defined in _add_sql_strings) """

        # Insert information into the error flags table
        self._add_sql_strings()
        for this_table, this_str in self.create_table_sql.items():
            self.execute_sql(this_str)
        error_data = [(0, '', 'No error'), (1, 'M', 'Improbable value flagged by QC'),
                        (2, 'N', 'Null Value'), (3, 'T', 'Value interpolated from adjacent values')]
        self.insert_into_table('error_flags', error_data)

    def insert_tide_file(self, file_list):
        """
        Add data from a set of files to the database.

        Parameters
        ----------
        file_list : list
            List of file names.

        """
        for this_file in file_list:
            print('Inserting data from file: ' + this_file)
            this_file_obj = bodc_annual_tide_file(this_file)
            try:
                site_id = self.select_qry('sites', "site_tla == '" + this_file_obj.site_tla + "'", 'site_id')[0][0]
            except:
                try:
                    current_id_max = np.max(self.select_qry('sites', None, 'site_id'))
                    site_id = int(current_id_max + 1)
                except:
                    site_id = 1

                site_data = [(site_id, this_file_obj.site_tla, this_file_obj.site_name, this_file_obj.lon, this_file_obj.lat, '')]
                self.debug_data = site_data
                self.insert_into_table('sites', site_data)

            site_id_list = [site_id] * len(this_file_obj.seconds_from_ref)
            table_data = list(zip(site_id_list, this_file_obj.seconds_from_ref, this_file_obj.elevation_data,
                            this_file_obj.elevation_flag, this_file_obj.residual_data, this_file_obj.residual_flag))
            self.insert_into_table('gauge_obs', table_data)

    def get_tidal_series(self, station_identifier, start_date_dt=None, end_date_dt=None):
        """
        Extract a time series of tidal elevations for a given station.

        Parameters
        ----------
        station_identifier : str
            Database station identifier.
        start_date_dt, end_date_dt : datetime.datetime, optional
            Give start and/or end times to extract from the database. If omitted, all data are returned.

        Returns
        -------
        dates : np.ndarray
            Array of datetime objects.
        data : np.ndarray
            Surface elevation and residuals from the database for the given station.

        """
        select_str = "time_int, elevation, elevation_flag"
        table_name = "gauge_obs as go"
        inner_join_str = "sites as st on st.site_id = go.site_id"

        if isinstance(station_identifier, str):
            where_str = "st.site_tla = '" + station_identifier + "'"
        else:
            where_str = "st.site_id = " + str(int(station_identifier))

        if start_date_dt is not None:
            start_sec = dt_to_epochsec(start_date_dt)
            where_str += " and go.time_int >= " + str(start_sec)
        if end_date_dt is not None:
            end_sec = dt_to_epochsec(end_date_dt)
            where_str += " and go.time_int <= " + str(end_sec)
        order_by_str = 'go.time_int'
        return_data = self.select_qry(table_name, where_str, select_str, order_by_str, inner_join_str)
        if not return_data:
            print('No data available')
            dates, data = None, None
        else:
            return_data = np.asarray(return_data)
            date_list = [epochsec_to_dt(this_time) for this_time in return_data[:,0]]
            dates, data = np.asarray(date_list), return_data[:, 1:]

        return dates, data

    def get_gauge_locations(self, long_names=False):
        """
        Extract locations and names of the tide gauges from the database.

        Parameters
        ----------
        long_names : bool, optional
            If True, return the 'nice' long names rather than the station identifiers.

        Returns
        -------
        tla_name : np.ndarray
            List of tide gauge names.
        lon_lat : np.ndarray
            Positions of the gauges.

        """
        gauge_site_data = np.asarray(self.select_qry('sites', None))
        if long_names:
            tla_name = gauge_site_data[:, 2]
        else:
            tla_name = gauge_site_data[:, 1]
        lon_lat = np.asarray(gauge_site_data[:, 3:5], dtype=float)

        return tla_name, lon_lat

    def get_nearest_gauge_id(self, lon, lat, threshold=np.inf):
        """
        Get the ID of the gauge closest to the position given by `lon' and `lat'.

        lon, lat : float
            Position for which to search for the nearest tide gauge.

        Returns
        -------
        closest_gauge_id : int
            Database ID for the gauge closest to `lon' and `lat'.
        min_dist : float
            Distance in metres between `lon' and `lat' and the gauge.
        threshold : float
            Threshold distance in metres (inclusive) within which gauges must be from the given position. If no
            gauges are found within this distance, the gauge ID is None.

        """

        sites_lat_lon = np.asarray(self.select_qry('sites', None, 'site_id, lat, lon'))
        min_dist = np.inf
        closest_gauge_id = None  # we should make this False or None or something
        for this_row in sites_lat_lon:
            this_dist = vincenty_distance([lat, lon], [this_row[1], this_row[2]])
            if this_dist < min_dist:
                min_dist = this_dist
                closest_gauge_id = this_row[0]
        if min_dist >= threshold:
            closest_gauge_id = None
        else:
            closest_gauge_id = int(closest_gauge_id)

        return closest_gauge_id, min_dist

    def _add_sql_strings(self):
        """ Function to define the database structure. """
        bodc_tables = {'gauge_obs': ['site_id integer NOT NULL', 'time_int integer NOT NULL',
                                     'elevation real NOT NULL', 'elevation_flag integer', 'residual real', 'residual_flag integer',
                                     'PRIMARY KEY (site_id, time_int)', 'FOREIGN KEY (site_id) REFERENCES sites(site_id)',
                                     'FOREIGN KEY (elevation_flag) REFERENCES error_flags(flag_id)',
                                     'FOREIGN KEY (residual_flag) REFERENCES error_flags(flag_id)'],
                       'sites': ['site_id integer NOT NULL', 'site_tla text NOT NULL', 'site_name text', 'lon real', 'lat real',
                                 'other_stuff text', 'PRIMARY KEY (site_id)'],
                       'error_flags': ['flag_id integer NOT NULL', 'flag_code text', 'flag_description text']}

        for this_key, this_val in bodc_tables.items():
            self.make_create_table_sql(this_key, this_val)

class bodc_annual_tide_file():
    def __init__(self, file_name, header_length=11):
        """
        Assumptions: file name of the form yearTLA.txt

        """
        bodc_annual_tide_file._clean_tide_file(file_name, header_length)
        with open(file_name) as f:
            header_lines = [next(f) for this_line in range(header_length)]

        for this_line in header_lines:
            if 'ongitude' in this_line:
                self.lon = [float(s) for s in this_line.split() if bodc_annual_tide_file._is_number(s)][0]
            if 'atitude' in this_line:
                self.lat = [float(s) for s in this_line.split() if bodc_annual_tide_file._is_number(s)][0]
            if 'Site' in this_line:
                site_str_raw = this_line.split()[1:]
                if len(site_str_raw) == 1:
                    site_str = site_str_raw[0]
                else:
                    site_str = ''
                    for this_str in site_str_raw:
                        site_str += this_str
        self.site_name = site_str
        self.site_tla = file_name.split('/')[-1][4:7]

        raw_data = np.loadtxt(file_name, skiprows=header_length, dtype=bytes).astype(str)

        seconds_from_ref = []
        for this_row in raw_data:
            this_dt_str = this_row[1] + ' ' + this_row[2]
            this_seconds_from_ref = dt_to_epochsec(dt.datetime.strptime(this_dt_str, '%Y/%m/%d %H:%M:%S'))
            seconds_from_ref.append(int(this_seconds_from_ref))
        self.seconds_from_ref = seconds_from_ref

        elevation_data = []
        elevation_flag = []
        residual_data = []
        residual_flag = []
        for this_row in raw_data:
            meas, error_code = bodc_annual_tide_file._parse_tide_obs(this_row[3])
            elevation_data.append(meas)
            elevation_flag.append(error_code)
            meas, error_code = bodc_annual_tide_file._parse_tide_obs(this_row[4])
            residual_data.append(meas)
            residual_flag.append(error_code)
        self.elevation_data = elevation_data
        self.elevation_flag = elevation_flag
        self.residual_data = residual_data
        self.residual_flag = residual_flag

    @staticmethod
    def _parse_tide_obs(in_str):
        error_code_dict = {'M':1, 'N':2, 'T':3}
        try:
            int(in_str[-1])
            error_code = 0
            meas = float(in_str)
        except:
            error_code_str = in_str[-1]
            meas = float(in_str[0:-1])
            try:
                error_code = error_code_dict[error_code_str]
            except:
                print('Unrecognised error code')
                return
        return meas, error_code

    @staticmethod
    def _is_number(s):
        try:
            float(s)
            return True
        except ValueError:
            return False

    @staticmethod
    def _clean_tide_file(file_name, header_length):
        sed_str = "sed -i '"+ str(header_length + 1) + ",$ {/^ *[0-9]/!d}' " + file_name
        sp.call([sed_str], shell=True)

#################################################################################################################
"""

Validation against L4 and E1 CTD and buoy data

observations_meta_data = {'buoy_name':'E1', 'year':'2006', 'ctd_new_file_type': False,
            'ctd_datadir':'/data/euryale4/backup/mbe/Data/WCO_data/E1/CTD_data/2006',
            'buoy_filepath':None, 'lon':-4.368, 'lat':50.035}

observations_meta_data = {'buoy_name':'L4', 'year':'2015', 'ctd_new_file_type': True, 'ctd_filepath':'./data/e1_data_2015.txt',
            'buoy_filepath': , '/data/euryale4/backup/mbe/Data/WCO_data/L4/Buoy_data/l4_cont_data_2015.txt', 'lon':-4.217, 'lat':50.250}

model_filestr_lambda = lambda m: '/data/euryale4/backup/mbe/Models/FVCOM/tamar_v2/run/output/depth_tweak2/2006/{:02d}/tamar_v2_0001.nc'.format(m)
available_months = np.arange(1,13)
model_file_list = [model_filestr_lambda(this_month) for this_month in available_months]

"""

class db_wco(validation_db):
    """ Work with an SQL database of data from PML's Western Channel Observatory. """

    def make_wco_tables(self):
        """
        Make the complete set of empty tables for data to be inserted into (as defined in _add_sql_strings).

        """

        # Insert information into the error flags table
        self._add_sql_strings()
        for this_table, this_str in self.create_table_sql.items():
            self.execute_sql(this_str)
        sites_data = [(0, 'L4',-4.217,50.250, ' '), (1, 'E1',-4.368,50.035,' ')]
        self.insert_into_table('sites', sites_data)
        measurement_type_data = [(0,'CTD measurements'), (1, 'Surface buoy measurements')]
        self.insert_into_table('measurement_types', measurement_type_data)
        self.execute_sql('create index date_index on obs (time_int);')

    def _add_sql_strings(self):
        wco_tables = {'obs':['buoy_id integer NOT NULL', 'time_int integer NOT NULL',
                                    'depth real NOT NULL', 'temp real', 'salinity real', 'measurement_flag integer NOT NULL',
                                    'PRIMARY KEY (buoy_id, depth, measurement_flag, time_int)', 'FOREIGN KEY (buoy_id) REFERENCES sites(buoy_id)',
                                    'FOREIGN KEY (measurement_flag) REFERENCES measurement_types(measurement_flag)'],
                       'sites':['buoy_id integer NOT NULL', 'buoy_name text', 'lon real', 'lat real',
                                    'other_stuff text', 'PRIMARY KEY (buoy_id)'],
                       'measurement_types':['measurement_flag integer NOT NULL', 'measurement_description text', 'PRIMARY KEY (measurement_flag)']}

        for this_key, this_val in wco_tables.items():
            self.make_create_table_sql(this_key, this_val)

    def insert_CTD_file(self, filestr, buoy_id):
        file_obj = WCO_obs_file(filestr)
        self._insert_obs(file_obj, buoy_id, 0.0)

    def insert_buoy_file(self, filestr, buoy_id):
        file_obj = WCO_obs_file(filestr, depth=0)
        self._insert_obs(file_obj, buoy_id, 1.0)

    def insert_CTD_dir(self, dirstr, buoy_id):
        file_obj = CTD_dir(dirstr)
        self._insert_obs(file_obj, buoy_id, 0.0)

    def _insert_obs(self, file_obj, buoy_id, measurement_id):
        epoch_sec_timelist = []
        for this_time in file_obj.observation_dict['dt_time']:
            epoch_sec_timelist.append(dt_to_epochsec(this_time))
        buoy_id_list = np.tile(buoy_id, len(epoch_sec_timelist))
        measurement_id_list = np.tile(measurement_id, len(epoch_sec_timelist))

        table_data = list(zip(buoy_id_list, epoch_sec_timelist, file_obj.observation_dict['depth'], file_obj.observation_dict['temp'],
                        file_obj.observation_dict['salinity'], measurement_id_list))
        self.insert_into_table('obs', table_data)

    def get_observations(self, buoy_name, start_date_dt=None, end_date_dt=None, measurement_id=None):
        select_str = "time_int, depth, temp, salinity"
        table_name = "obs as go"
        inner_join_str = "sites as st on st.buoy_id = go.buoy_id"

        where_str = "st.buoy_name = '" + buoy_name + "'"

        if start_date_dt is not None:
            start_sec = dt_to_epochsec(start_date_dt)
            where_str += " and go.time_int >= " + str(start_sec)
        if end_date_dt is not None:
            end_sec = dt_to_epochsec(end_date_dt)
            where_str += " and go.time_int <= " + str(end_sec)
        order_by_str = 'go.time_int, go.depth'

        return_data = self.select_qry(table_name, where_str, select_str, order_by_str, inner_join_str)

        if not return_data:
            dates, data = None, None
            print('No data available')
        else:
            return_data = np.asarray(return_data)
            date_list = [epochsec_to_dt(this_time) for this_time in return_data[:,0]]
            dates, data = np.asarray(date_list), return_data[:,1:]

        return dates, data

class WCO_obs_file():
    def __init__(self, filename, depth=None):
        self._setup_possible_vars()
        self.observation_dict = self._add_file(filename)
        if depth is not None:
            self.observation_dict['depth'] = np.tile(depth, len(self.observation_dict['dt_time']))

    def _add_file(self,filename,remove_undated=True):
        # remove duff lines
        sed_str = "sed '/^-9.990e/d' " + filename + " > temp_file.txt"
        sp.call(sed_str, shell=True)

        # some files have multiple records of differing types...helpful
        temp_str = 'YKA123ASD'
        file_split_str = '''awk '/^[^0-9]/{g++} { print $0 > "''' + temp_str + '''"g".txt"}' temp_file.txt'''
        sp.call(file_split_str, shell=True)
        temp_file_list = gb.glob(temp_str + '*')

        obs_dict_list = []
        for this_file in temp_file_list:
            this_obs = self._add_file_part(this_file)
            if not remove_undated or 'dt_time' in this_obs:
                obs_dict_list.append(this_obs)
        rm_file = [os.remove(this_file) for this_file in temp_file_list]
        return {this_key:np.hstack([this_dict[this_key] for this_dict in obs_dict_list]) for this_key in obs_dict_list[0]}

    def _add_file_part(self, filename):
        # seperate header and clean out non numeric lines
        head_str ="head -1 " + filename + " > temp_header_file.txt"
        sed_str = "sed '/^[!0-9]/!d' " + filename + " > temp_file.txt"
        sp.call(head_str, shell=True)
        sp.call(sed_str, shell=True)

        # Load the files, some use semi-colon delimiters, some whitespace...
        if ';' in str(np.loadtxt('temp_header_file.txt', delimiter='no_delimination_needed', dtype=str)):
            observations_raw = np.loadtxt('temp_file.txt', delimiter=';',dtype=str)
            observations_header = np.loadtxt('temp_header_file.txt', delimiter=';',dtype=str)
        elif ',' in str(np.loadtxt('temp_header_file.txt', delimiter='no_delimination_needed', dtype=str)):
            observations_raw = np.loadtxt('temp_file.txt', delimiter=',',dtype=str)
            observations_header = np.loadtxt('temp_header_file.txt', delimiter=',',dtype=str)
        else:
            observations_raw = np.loadtxt('temp_file.txt', dtype=str)
            observations_header = np.loadtxt('temp_header_file.txt', dtype=str)

        # Clean up temp files
        os.remove('temp_file.txt')
        os.remove('temp_header_file.txt')

        # Find the relevant columns and pull out temp, salinity, date, etc if available
        observation_dict = {}
        time_vars = []

        for this_var, this_possible in self.possible_vars.items():
            if np.any(np.isin(this_possible, observations_header)):
                this_col = np.where(np.isin(observations_header, this_possible))[0]
                if this_var == 'time' or this_var =='date' or this_var=='Jd':
                    observation_dict[this_var] = np.squeeze(np.asarray(observations_raw[:,this_col], dtype=str))
                    time_vars.append(this_possible[np.isin(this_possible, observations_header)])
                else:
                    observation_dict[this_var] = np.squeeze(np.asarray(observations_raw[:,this_col], dtype=float))
        if 'date' in observation_dict:
            observation_dict['dt_time'] = self._parse_dates_to_dt(observation_dict, time_vars)
        return observation_dict

    def _setup_possible_vars(self):
        self.possible_vars = {'temp':np.asarray(['Tv290C', 'SST', ' Mean SST (degC)']), 'salinity':np.asarray(['Sal00', 'Sal', ' Mean SST (degC)']),
                        'depth':np.asarray(['DepSM']),    'date':np.asarray(['mm/dd/yyyy', 'Year', ' Date (YYMMDD)']),
                        'julian_day':np.asarray(['Jd']), 'time':np.asarray(['hh:mm:ss', 'Time', ' Time (HHMMSS)'])}

    @staticmethod
    def _parse_dates_to_dt(obs_dict, time_vars):
        dt_list = []
        if np.any(np.isin('mm/dd/yyyy', time_vars)):
            for this_time, this_date in zip(obs_dict['time'], obs_dict['date']):
                dt_list.append(dt.datetime.strptime(this_date + ' ' + this_time, '%m/%d/%Y %H:%M:%S'))
        elif np.any(np.isin('Year', time_vars)):
            for this_time, (this_jd, this_year) in zip(obs_dict['time'], zip(obs_dict['julian_day'], obs_dict['date'])):
                dt_list.append(dt.datetime(int(this_year),1,1) + dt.timedelta(days=int(this_jd) -1) +
                                dt.timedelta(hours=int(this_time.split('.')[0])) + dt.timedelta(minutes=int(this_time.split('.')[1])))
        elif np.any(np.isin(' Date (YYMMDD)', time_vars)):
            for this_time, this_date in zip(obs_dict['time'], obs_dict['date']):
                dt_list.append(dt.datetime.strptime(this_date + ' ' + this_time, '%y%m%d %H%M%S'))
        else:
            print('Date parser not up to date with possible vars')
            dt_list = None

        return np.asarray(dt_list)

class CTD_dir(WCO_obs_file):
    def __init__(self, dirname):
        all_files = os.listdir(dirname)
        dt_list = []
        observation_dict_list = []
        self._setup_possible_vars()

        for this_file in all_files:
            print('Processing file {}'.format(this_file))
            try:
                observation_dict_list.append(self._add_file(dirname + this_file, remove_undated=False))
                date_str = '20' + this_file[0:2] + '-' + this_file[2:4] + '-' + this_file[4:6]
                this_dt = dt.datetime.strptime(date_str, '%Y-%m-%d') + dt.timedelta(hours=12)
                dt_list.append(np.tile(this_dt, len(observation_dict_list[-1]['temp'])))

            except ValueError:
                print('Error in file {}'.format(this_file))
        # Flatten the list of dictionaries to one dictionary
        self.observation_dict = {this_key:np.hstack([this_dict[this_key] for this_dict in observation_dict_list]) for this_key in observation_dict_list[0]}
        self.observation_dict['dt_time'] = np.hstack(dt_list)

"""
Do the comparison

"""

class comp_data():
    def __init__(self, buoy_list, file_list_or_probe_dir, wco_database, max_time_threshold=dt.timedelta(hours=1), max_depth_threshold = 100, probe_depths=None):
        self.file_list_or_probe_dir
        self.database_obj = wco_database
        self.buoy_list = buoy_list

        self.model_data = {}
        if probe_depths is not None:
            for this_ind, this_buoy in enumerate(buoy_list):
                self.model_data[this_buoy]['depth'] = probe_depths[this_ind]

        self.observations = {}
        for this_buoy in self.buoy_list:
            self.observations[this_buoy] = {}

    def retrieve_file_data(self):
        pass

    def retrieve_obs_data(self, buoy_name, var, measurement_type=None):
        if not hasattr(self.model_date_mm):
            print('Retrieve model data first')
            return
        obs_dt, obs_raw = self.database_obj.get_obs_data(buoy_name, var, self.model_date_mm[0], self.model_date_mm[1], measurement_type)
        obs_depth = obs_raw[:,0]
        obs_var = obs_raw[:,1]
        self.observations[buoy_name][var] = obs_dict

    def get_comp_data_interpolated(self, buoy_name, var_list):
        if not hasattr(self, comp_dict):
            self.comp_dict = {}

        obs_dates = np.unique([this_date.date() for this_date in observations['time']])
        obs_comp = {}
        for this_ind, this_obs_time in enumerate(obs_dates):
            if this_obs_date >= model_time_mm[0].date() and this_obs_date <= model_time_mm[1].date():
                this_obs_choose = [this_time.date() == this_obs_date for this_time in self.observations[buoy_name]['dt_time']]
                t
                this_time_before, this_time_after = self.model_closest_both_times(this_obs_time)

                this_obs_deps = self.observations[buoy_name]['depth'][this_obs_choose]
                for var in var_list:
                    this_obs = self.observations[buoy_name][var][this_obs_choose]
                    this_model = np.squeeze(fvcom_data_reader.data.temp[this_time_close,...])
                    this_model_interp = np.squeeze(np.interp(this_obs_deps, model_depths, this_model))

                    try:
                        obs_comp[var].append(this_comp)
                    except KeyError:
                        obs_comp[var] = this_comp

                    max_obs_depth.append(np.max(this_obs_deps))

        self.comp_dict[buoy_name] = obs_comp

    def comp_data_nearest(self, buoy_name, var_list):
        pass

    def model_closest_time():
        pass

class comp_data_filereader(comp_data):
    def retrieve_file_data(self):
        where_str = 'buoy_name in ('
        for this_buoy in self.buoy_list:
            where_str+=this_buoy + ','
        where_str = where_str[0:-1] + ')'
        buoy_lat_lons = self.wco_database.select_query('sites', where_str, 'buoy_name, lon, lat')

        first_file = True
        model_all_dicts = {}

        for this_file in self.file_list:
            if first_file:
                fvcom_data_reader = FileReader(this_file, ['temp', 'salinity'])
                close_node = []
                for this_buoy_ll in buoy_lat_lons:
                    close_node.append()

                model_depths = fvcom_data_reader.grid.siglay * fvcom_data_reader.grid.h * -1

                for this_buoy in self.buoy_list:
                    model_all_dicts[this_buoy] = {'dt_time': mod_times, 'temp': mod_t_vals, 'salinity': mod_s_vals}

                first_file = False
            else:
                fvcom_data_reader = FileReader(this_file, ['temp', 'salinity'], dims={'node':[close_node]})
                for this_buoy in self.buoy_list:
                    model_all_dicts

        model_depths = fvcom_data_reader.grid.siglay * fvcom_data_reader.grid.h * -1
        for this_buoy in self.buoy_list:
            model_all_dicts[this_buoy]['depth'] = model_dp

        for this_buoy in self.buoy_list:
            self.model_data[this_buoy] = model_all_dicts[this_buoy]

    def model_closest_time(self, find_time):
        return closest_time

class comp_data_probe(comp_data):
    def retrieve_file_data():
        for this_buoy in self.buoy_list:
            t_filelist = []
            s_filelist = []
            for this_dir in self.file_or_probe_dir_list:
                t_filelist.append(this_dir + this_buoy + '_t1.dat')
                s_filelist.append(this_dir + this_buoy + '_s1.dat')
            mod_times, mod_t_vals, mod_pos = pf.read.read_probes(t_filelist, locations=True, datetimes=True)
            mod_times, mod_s_vals, mod_pos = pf.read.read_probes(s_filelist, locations=True, datetimes=True)
            model_dict = {'dt_time':mod_times, 'temp':mod_t_vals, 'salinity':mod_s_vals}
            self.model_data[this_buoy] = model_dict

def wco_model_comparison(model_file_list, obs_database_file):

    temp_comp = []
    sal_comp = []
    dates_comp = []
    max_obs_depth = []

    obs_dates = np.unique([this_date.date() for this_date in observations['time']])
    for this_ind, this_obs_date in enumerate(obs_dates):
        if this_obs_date >= model_time_mm[0].date() and this_obs_date <= model_time_mm[1].date():
            this_obs_choose = [this_time.date() == this_obs_date for this_time in observations['time']]
            this_obs_time = np.min(observations['time'][this_obs_choose])
            this_time_close = fvcom_data_reader.closest_time(this_obs_time)

            this_obs_deps = observations['h'][this_obs_choose]
            this_obs_temp = observations['temp'][this_obs_choose]
            this_obs_salinity = observations['salinity'][this_obs_choose]

            this_obs_temp_interp = np.squeeze(np.interp(model_depths, this_obs_deps, this_obs_temp))
            this_obs_salinity_interp = np.squeeze(np.interp(model_depths, this_obs_deps, this_obs_salinity))

            this_model_temp = np.squeeze(fvcom_data_reader.data.temp[this_time_close,...])
            this_model_salinity = np.squeeze(fvcom_data_reader.data.salinity[this_time_close,...])

            temp_comp.append(np.asarray([this_obs_temp_interp, this_model_temp]))
            sal_comp.append(np.asarray([this_obs_salinity_interp, this_model_salinity]))
            dates_comp.append(this_obs_time)
            max_obs_depth.append(np.max(this_obs_deps))

    ctd_comp = {'temp':temp_comp, 'salinity':sal_comp, 'dates':dates_comp, 'max_depth':max_obs_depth}

    observations = get_buoy_obs(obs_meta_data)

    if observations:
        buoy_temp_comp = []
        buoy_sal_comp = []
        buoy_dates_comp = []

        for this_ind, this_obs_time in enumerate(observations['time']):
            if this_obs_time >= model_time_mm[0] and this_obs_time <= model_time_mm[1]:
                this_time_close = fvcom_data_reader.closest_time(this_obs_time)
                buoy_temp_comp.append([observations['temp'][this_ind], fvcom_data_reader.data.temp[this_time_close,0]])
                buoy_sal_comp.append([observations['salinity'][this_ind], fvcom_data_reader.data.salinity[this_time_close,0]])
                buoy_dates_comp.append([this_obs_time, this_time_close])

        buoy_comp = {'temp':buoy_temp_comp, 'salinity':buoy_sal_comp, 'dates':buoy_dates_comp}

    else:
        buoy_comp = {}

    return ctd_comp, buoy_comp


import argparse
import time
import re
import os
import datetime

import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data import Dataset, DataLoader, random_split
from torch.nn.utils.rnn import pack_padded_sequence, pad_sequence

from vocab import Vocab

class Merger(nn.Module):
    def __init__(self, args, vocab):
        super().__init__()
        self.bigram_encoder = Merger.BigramCharRNN(
            args,
            output_vocab_size=2,
            input_vocab=vocab.src.char2id)

    def forward(self, bigrams, lengths_bigram):
        outputs_bigram = self.bigram_encoder(bigrams, lengths_bigram)
        return outputs_bigram

    class BigramCharRNN(nn.Module):
        def __init__(self,
                    args,
                    output_vocab_size,
                    input_vocab):
            super().__init__()
            self.input_vocab = input_vocab
            self.hidden_dim = args.rnn_dim_char
            self.num_layers = args.rnn_layers
            self.word_embeddings = nn.Embedding(len(input_vocab), args.ce_dim)
            self.lstm = nn.LSTM(args.ce_dim,
                            args.rnn_dim_char,
                            num_layers=args.rnn_layers,
                            bidirectional=True)
            self.fc0 = nn.Linear(args.rnn_dim_char, 64)
            self.fc1 = nn.Linear(64, output_vocab_size)
            self.relu = nn.ReLU()
            self.dropout = nn.Dropout(0.1)

        def forward(self, bigrams, lengths):
            embeds = self.dropout(self.word_embeddings(bigrams))
            packed_embeds = pack_padded_sequence(
                embeds, lengths, enforce_sorted=False)
            _, hidden = self.lstm(packed_embeds)
            hidden = sum(hidden[i][j, :, :]
                            for i in range(2) for j in range(self.num_layers))
            tag_space = self.relu(self.fc0(hidden))
            tag_space = self.fc1(tag_space)
            return tag_space

class Network:
    def __init__(self, args) -> None:
        self.device = torch.device(
            f'cuda:{args.gpu_index}' if torch.cuda.is_available() else 'cpu')
        self.vocab = Vocab.load(args.vocab_path)
        self.load_data(args)
        self.model = Merger(args, vocab=self.vocab).to(self.device)
        self.loss_function = nn.CrossEntropyLoss()
        self.optimizer = optim.Adam(self.model.parameters())
        self.scheduler = ReduceLROnPlateau(self.optimizer, 'min', patience=3)

    class BigramData(Dataset):
        def __init__(self, data) -> None:
            self.data = data

        def __getitem__(self, index):
            src = torch.tensor(self.data[index][0], dtype=torch.long)
            label = torch.tensor(self.data[index][1], dtype=torch.long)
            return src, label

        def __len__(self):
            return len(self.data)

    def load_data(self, args, train_split=0.8):
        with open(args.data) as f:
            data = [line.split('\t') for line in f.readlines()]
            src = self.vocab.src.words2indices([list(line[0].strip()) for line in data])
            labels = [0 if bigram[1].strip() == '0' else 1 for bigram in data]
            src, labels = src[:args.data_size], labels[:args.data_size]
            data = list(zip(src, labels))
            lengths = [int(len(src)*train_split), int(len(src)*(1-train_split))]
            if sum(lengths) != len(src):
                lengths[0] += len(src) - sum(lengths)
            train_data, dev_data = random_split(data, lengths)

            train_data = Network.BigramData(train_data)
            dev_data = Network.BigramData(dev_data)

        def generate_batch(data_batch):
            src_batch, labels_batch = [], []
            lengths_bigram = []
            for src_item, label_item in data_batch:
                lengths_bigram.append(len(src_item))
                src_batch.append(src_item)
                labels_batch.append(label_item)

            src_batch = pad_sequence(
                src_batch, padding_value=self.vocab.src.word2id['<pad>'])
            labels_batch = torch.stack(labels_batch)
            return (src_batch, labels_batch), lengths_bigram

        self.train_iter = DataLoader(train_data, batch_size=args.batch_size,
                                     shuffle=True, collate_fn=generate_batch)
        self.dev_iter = DataLoader(train_data, batch_size=args.batch_size,
                                     collate_fn=generate_batch)

    @staticmethod
    def epoch_time(start_time: int,
                   end_time: int):
        elapsed_time = end_time - start_time
        elapsed_mins = int(elapsed_time / 60)
        elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
        return elapsed_mins, elapsed_secs

    def train(self, args):
        metrics = {'train_loss': [], 'dev_loss': [], 'dev_acc': []}
        for epoch in range(args.epochs):
            self.model.train()
            print(f'Epoch {epoch+1}/{args.epochs}')
            epoch_loss = 0
            start_time = time.time()
            for iteration, train_batch in enumerate(self.train_iter):
                (bigrams, labels), lengths = train_batch
                self.model.zero_grad()
                bigrams, labels = bigrams.to(self.device), labels.to(self.device)
                tag_scores = self.model(bigrams, lengths)
                tag_scores = tag_scores.view(-1, tag_scores.shape[-1])
                labels = labels.view(-1)
                loss = self.loss_function(tag_scores, labels)
                loss.backward()
                self.optimizer.step()
                epoch_loss += loss.item()
                if iteration and iteration % 100 == 0 and len(self.train_iter) - iteration > 10 \
                        or iteration + 1 == len(self.train_iter):
                    for param_group in self.optimizer.param_groups:
                        lr = param_group['lr']
                    print(
                        f'Batch {iteration}/{len(self.train_iter)-1}\t| Loss {loss.item():.7f} | lr {lr}')
                    metrics['train_loss'].append(epoch_loss / iteration)
            end_time = time.time()
            epoch_mins, epoch_secs = Network.epoch_time(start_time, end_time)
            val_metrics = self.evaluate()
            for m in metrics:
                if m != 'train_loss':
                    metrics[m].append(val_metrics[m])
            print(
                f'Epoch {epoch+1}/{args.epochs} | Time: {epoch_mins}m {epoch_secs}s')
            print(
                f"\tTrain Loss: {metrics['train_loss'][-1]:.7f} | Dev. Loss: {metrics['dev_loss'][-1]:.7f} | Dev. Acc.: {metrics['dev_acc'][-1]:.1%}")
            print()
            self.scheduler.step(metrics['dev_loss'][-1])

        self.save_model(args)
        return metrics

    def evaluate(self):
        self.model.eval()
        with torch.no_grad():
            correct, total = 0, 0
            epoch_loss = 0
            for (bigrams, labels), lengths in self.dev_iter:
                # Loss
                bigrams, labels = bigrams.to(self.device), labels.to(self.device)
                output = self.model(bigrams, lengths)
                loss = self.loss_function(output, labels)
                epoch_loss += loss.item()

                # Accuracy
                output = output.argmax(-1)
                correct += torch.sum(output == labels)
                total += labels.shape[0]
        metrics = {}
        metrics['dev_acc'] = correct / total
        metrics['dev_loss'] = epoch_loss / len(self.dev_iter)
        return metrics

    def predict(self):
        self.model.eval()
        pred, gold = [], []
        inputs_bigram = []
        with torch.no_grad():
            for (bigrams, labels), lengths in self.dev_iter:
                bigrams, labels = bigrams.to(self.device), labels.to(self.device)
                output = self.model(bigrams, lengths)
                output = output.argmax(-1)
                pred += list(output.detach().cpu().numpy())
                gold += list(labels.detach().cpu().numpy())
                inputs_bigram += list(bigrams.permute(1, 0).detach().cpu().numpy())

        return inputs_bigram, gold, pred

    @staticmethod
    def load_model(model_path: str):
        params = torch.load(model_path)
        args = params['args']
        network = Network(args)
        network.model.load_state_dict(params['state_dict'])
        return network

    def save_model(self, args):
        save_path = os.path.join(args.cpt, args.logdir) + '.pt'
        print('Saving model parameters to [%s]\n' % save_path)
        params = {
            'args': args,
            'state_dict': self.model.state_dict()
        }
        torch.save(params, save_path)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch_size", default=1048,
                        type=int, help="Batch size.")
    parser.add_argument("--epochs", default=50, type=int,
                        help="Number of epochs.")
    parser.add_argument("--ce_dim", default=32, type=int,
                        help="Word embedding dimension.")
    parser.add_argument("--we_dim", default=64, type=int,
                        help="Word embedding dimension.")
    parser.add_argument("--rnn_dim_char", default=64,
                        type=int, help="RNN cell dimension.")
    parser.add_argument("--rnn_dim", default=128,
                        type=int, help="RNN cell dimension.")
    parser.add_argument("--rnn_layers", default=1,
                        type=int, help="Number of RNN layers.")
    parser.add_argument("--data_size", default=10000, type=int,
                        help="Maximum number of examples to load.")
    parser.add_argument("--gpu_index", default=0, type=int,
                        help="Index of GPU to be used.")
    parser.add_argument("--vocab", dest='vocab_path',
                        default="/Users/chriscay/Library/Mobile Documents/com~apple~CloudDocs/NYUAD/camel_morph/sandbox_files/root_generator/data_nn_vocab.json", type=str,
                        help="Path to vocab JSON file.")
    parser.add_argument("--data",
                        default="/Users/chriscay/Library/Mobile Documents/com~apple~CloudDocs/NYUAD/camel_morph/sandbox_files/root_generator/data_nn.tsv", type=str,
                        help="Path to file with bigrams dataset.")
    parser.add_argument("--cpt", default='/Users/chriscay/Library/Mobile Documents/com~apple~CloudDocs/NYUAD/camel_morph/sandbox_files/root_generator/model_weights', type=str,
                        help="Directory to save the model checkpoints to.")
    parser.add_argument("--logs", default='/Users/chriscay/Library/Mobile Documents/com~apple~CloudDocs/NYUAD/camel_morph/sandbox_files/root_generator/logs', type=str,
                        help="Directory to save the model checkpoints to.")
    parser.add_argument("--load", default='', type=str,
                        help="Directory to save the model checkpoints to.")

    parser.add_argument("--seed", default=42, type=int, help="Random seed.")
    args = parser.parse_args([] if "__file__" not in globals() else None)

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    args.logdir = "{}-{}-{}".format(
        os.path.basename(globals().get("__file__", "notebook").split('.')[0]),
        datetime.datetime.now().strftime("%Y-%m-%d_%H:%M:%S"),
        ",".join(("{}={}".format(re.sub(
            "(.)[^_]*_?", r"\1", key), value) for key, value in sorted(vars(args).items()) if isinstance(value, int)))
    )

    if not args.load:
        network = Network(args)
        metrics = network.train(args)
        print(metrics)
    else:
        network = Network.load(args.load)
        inputs, gold, pred = network.predict()
        with open(os.path.join(args.logs, args.logdir), 'w') as f:
            for i, result in enumerate(pred):
                bigram = ''.join(
                    list(map(lambda x: network.vocab.src.id2char[x], inputs[0][i])))
                print(re.sub(r'(<pad>)+?', r'', bigram), file=f, end=' | ')
                print(result, 'gold:', gold[i], file=f, end= ' | ')

if __name__ == '__main__':
    main()


#!/usr/bin/env python3

"""Work Log
Record work activities and store to a sqlite database

Created: 2018
Last Update: 2018-06-05
Author: Alex Koumparos
"""
import datetime
import re

# from csv_manager import CsvManager
from db_manager import DBManager
import wl_settings as settings

class Menu:
    """The user-facing class that handles all interaction with the user and
    interfaces with the database manager.
    """
    # STATUS VARIABLES
    # ----------------
    quit = False

    # INITIALIZER
    # -----------
    def __init__(self, load_menu=True):
        """Instantiates the app, applies default settings and launches the
        main menu.
        """
        print("\nWORK LOG")
        print("========")
        self.OPTIONS = {
            'date format': settings.DATE_FORMATS['iso 8601'],
            'save format (date)': settings.DATE_FORMATS['iso 8601'],
            'case sensitive search': False,
            'entries per page': 10,
            'allow future dates': False,
            'earliest allowed date': datetime.datetime(1900, 1, 1),
        }
        self.current_record = 0
        self.current_page_start = 0
        if load_menu:
            menu = self.main_menu()
            while not self.quit:
                menu = menu()

    # MENU METHODS
    # ------------
    def main_menu(self):
        """This is the root menu. The user selects which activity to perform
        and then the method returns the function for the activity.
        """
        inputs = {
            'a': {'text': 'Add new entry',
                  'function': self.add_entry},
            's': {'text': 'Search in existing entries',
                  'function': self.search_entries},
            'o': {'text': 'Options',
                  'function': self.options},
            'q': {'text': 'Quit program',
                  'function': self.quit_program}
        }
        while True:
            print("\nMAIN MENU")
            print("What would you like to do?")
            for key, value in inputs.items():
                print("{}) {}".format(key, value['text']))
            user_entry = input("> ").lower()
            print(user_entry)
            if user_entry not in inputs.keys():
                continue
            return inputs[user_entry]['function']

    def add_entry(self):
        """This is the menu where the user can add a task that was completed
        """
        while True:
            print("\nADD ENTRY")
            print("Username")
            input_text = input("Enter your name > ")
            username = input_text
            date = None
            while date is None:
                print("Date of the Task")
                user_entry = self.date_entry()
                if user_entry[0] is not None:  # error
                    print(user_entry[0])
                    continue
                else:
                    date = user_entry[1]
                    date_string = self.date_to_string(date, target='file')
            print("Name of the Task")
            input_text = input("Enter the name of the task > ")
            task_name = input_text
            time_spent = None
            while time_spent is None:
                print("Time spent")
                print("Enter a whole number of minutes (rounded)")
                input_text = input("> ")
                try:
                    time_spent = int(input_text)
                except ValueError:
                    print("Invalid value, please try again")
                    continue
                if time_spent < 0:
                    print("Invalid value, please try again")
                    continue
            print("Notes")
            input_text = input("(Optional, leave blank for none) ")
            notes = input_text
            # call method to write data to file
            dbm = DBManager()
            file_data = {
                settings.HEADERS['user']: username,
                settings.HEADERS['date']: date_string,
                settings.HEADERS['task_name']: task_name,
                settings.HEADERS['duration']: time_spent,
                settings.HEADERS['notes']: notes
            }
            dbm.add_entry(file_data)
            return self.main_menu

    def options(self):
        """This is the menu where the user can specify user-configurable
        options
        """
        print('OPTIONS')

        print("Choose a display date format")
        menu_choices = list(settings.DATE_FORMATS.keys())
        menu_size = len(menu_choices)
        for i in range(len(menu_choices)):
            print("({}) - {}".format(i + 1, menu_choices[i]))
        input_text = input("> ")
        if input_text in [str(x) for x in range(1, menu_size + 1)]:
            choice = int(input_text) - 1
            choice = menu_choices[choice]
            print("You chose: {}".format(choice))
            self.OPTIONS['date format'] = settings.DATE_FORMATS[choice]
            print('going back to main menu')
        else:
            print("Invalid entry, returning to main menu")

        return self.main_menu

    def search_entries(self):
        """This is the search menu. The user selects how they want to search.
        """
        inputs = {
            'l': {'text': 'employee names List',
                  'function': self.search_employee},
            'e': {'text': 'Employee name Search',
                  'function': self.search_employee_text},
            'd': {'text': 'single Date',
                  'function': self.search_exact_date},
            'r': {'text': 'date Range',
                  'function': self.search_date_range},
            't': {'text': 'Time spent',
                  'function': self.search_time_spent},
            's': {'text': 'text Search',
                  'function': self.search_text_search},
            'b': {'text': 'Back to main menu',
                  'function': self.main_menu}
        }
        while True:
            print("\nSEARCH ENTRIES")
            print("How would you like to search?")
            for key, value in inputs.items():
                print("{}) {}".format(key, value['text']))
            user_entry = input("> ").lower()

            if user_entry not in inputs.keys():
                continue
            return inputs[user_entry]['function']

    def quit_program(self):
        print("Quitting")
        self.quit = True

    def present_results(self):
        """Show all the results from the search and then provide interaction
        choices
        """
        inputs = {
            'n': {'text': 'Next page',
                  'function': self.next_page},
            'p': {'text': 'Previous page',
                  'function': self.previous_page},
            'v': {'text': 'View detail',
                  'function': self.select_detail},
            'e': {'text': 'Edit',
                  'function': self.edit_record},
            'd': {'text': 'Delete',
                  'function': self.delete_record},
            'm': {'text': 'go back to Main menu',
                  'function': self.main_menu},
            'q': {'text': 'quit',
                  'function': self.quit_program},
        }
        if self.current_page_start == 0:
            del(inputs['p'])
        next_start = self.current_page_start + self.OPTIONS['entries per page']
        if next_start >= len(self.records):
            del(inputs['n'])
        print("\nSearch Results")
        if len(self.records) > next_start:
            current_page_end = next_start
        else:
            current_page_end = len(self.records) - 1
        for index in range(self.current_page_start, current_page_end + 1):
            value = self.records[index]
            short_form = self.display_entry(value, return_only=True)
            print("{}) {}".format(index + 1, short_form))

        print("\nAvailable actions:")
        for key, value in inputs.items():
            print('{}) {}'.format(key, value['text']))

        while True:
            user_entry = input("> ").lower()

            if user_entry not in inputs.keys():
                continue
            return inputs[user_entry]['function']

    def present_next_result(self):
        """Show the next available result"""
        inputs = {
            'p': {'text': 'Previous',
                  'function': self.previous_result},
            'n': {'text': 'Next',
                  'function': self.next_result},
            'b': {'text': 'Back to list view',
                  'function': self.present_results},
            'e': {'text': 'Edit',
                  'function': self.edit_current_record},
            'd': {'text': 'Delete',
                  'function': self.delete_current_record},
            'm': {'text': 'go back to Main menu',
                  'function': self.main_menu},
            'q': {'text': 'quit',
                  'function': self.quit_program},
        }
        if self.current_record == 0:
            del(inputs['p'])
        if self.current_record == len(self.records) - 1:
            del(inputs['n'])
        record = self.records[self.current_record]
        self.display_entry(record, verbose=True)

        print("\nAvailable actions:")
        for key, value in inputs.items():
            print('{}) {}'.format(key, value['text']))

        while True:
            user_entry = input("> ").lower()

            if user_entry not in inputs.keys():
                continue
            return inputs[user_entry]['function']

    # Specific Search Menus
    # .....................
    def search_employee(self):
        """This is the menu where the user is given a list of all employees
        who have entries, and can select a particular employee to see
        all their entries
        """
        print("\nSEARCH BY EMPLOYEE")
        # load the db manager
        dbm = DBManager()
        employee_names = dbm.view_employees()
        for i, value in enumerate(employee_names):
            print("{}) {}".format(i + 1, value['name']))
        selected_employee = None
        while selected_employee is None:
            user_input = input("> ")
            # perform input validation
            try:
                user_input = int(user_input) - 1
            except ValueError:
                print("Invalid value, try again")
                continue
            if user_input < 0:
                print("Value out of range. Try again.")
                continue
            try:
                selected_employee = employee_names[user_input]['name']
            except IndexError:
                print("Value out of range. Try again.")
                continue
            # when an employee is selected, show all the entries with that e'ee
            matching_records = dbm.view_everything(employee=selected_employee)
        self.records = matching_records
        self.current_record = 0
        return self.present_next_result

    def search_employee_text(self):
        """This is the menu where the user enters a text string and is presented
        with all employee names containing that string
        """
        print('FIND EMPLOYEE NAME USING TEXT STRING')
        print("Enter the text string to search on")
        input_text = input("> ")
        text_string = input_text
        # load db
        dbm = DBManager()
        employee_names = dbm.view_names_with_text(text_string)
        for i, value in enumerate(employee_names):
            print("{}) {}".format(i + 1, value['name']))
        selected_employee = None
        while selected_employee is None:
            user_input = input("> ")
            # perform input validation
            try:
                user_input = int(user_input) - 1
            except ValueError:
                print("Invalid value, try again")
                continue
            if user_input < 0:
                print("Value out of range. Try again.")
                continue
            try:
                selected_employee = employee_names[user_input]['name']
            except IndexError:
                print("Value out of range. Try again.")
                continue
            # when an employee is selected, show all the entries with that e'ee
            matching_records = dbm.view_everything(employee=selected_employee)
        self.records = matching_records
        self.current_record = 0
        return self.present_next_result

    def search_exact_date(self):
        """This is the menu where the user browses dates and entries and picks
        the date from a list
        """
        print("\nSEARCH EXACT DATE")
        # load the db manager
        dbm = DBManager()
        date_records = dbm.view_dates()
        for i, value in enumerate(date_records):
            value = self.date_to_string(value['date'])
            print("{}) {}".format(i + 1, value))
        selected_date = None
        while selected_date is None:
            user_input = input("> ")
            # perform input validation
            try:
                user_input = int(user_input) - 1
            except ValueError:
                print("Invalid value, try again")
                continue
            if user_input < 0:
                print("Value out of range. Try again.")
                continue
            try:
                selected_date = date_records[user_input]['date']
            except IndexError:
                print("Value out of range. Try again.")
                continue

            # when a date is selected, show all the entries with that date
            matching_records = dbm.view_entries_for_date(selected_date)
        self.records = matching_records
        self.current_record = 0
        return self.present_next_result

    def search_date_range(self):
        """This is the menu where the user can enter a from date and to date
        and get back every entry from within that range
        """
        print('SEARCH DATE RANGE')
        start_date = None
        end_date = None
        # get start_date
        while start_date is None:
            print("Start Date:")
            user_entry = self.date_entry()
            if user_entry[0] is not None:  # error
                print(user_entry[0])
                continue
            else:
                start_date = user_entry[1]
        # get end_date
        while end_date is None:
            print("End Date:")
            user_entry = self.date_entry()
            if user_entry[0] is not None:  # error
                print(user_entry[0])
                continue
            else:
                end_date = user_entry[1]
        # load db
        dbm = DBManager()
        # switch start and end dates if end < start
        if end_date < start_date:
            current_date = end_date
            end_date = start_date
            start_date = end_date
        else:
            current_date = start_date
        # get all records in date range
        matching_records = dbm.view_entries_for_date_range(start_date,
                                                           end_date)

        print("\nShowing entries:")
        if len(matching_records) == 0:
            print("\nNo matches, returning to search menu")
            return self.search_entries
        self.records = matching_records
        self.current_record = 0
        return self.present_next_result

    def search_time_spent(self):
        """This is the menu where the user enters the number of minutes a task
        took and be able to choose one to see entries from
        """
        print('SEARCH BY TIME SPENT')
        print("Time spent")
        time_spent = None
        while time_spent is None:
            input_text = input("Enter a whole number of minutes (rounded) ")
            try:
                time_spent = int(input_text)
            except ValueError:
                print("Invalid value")
                continue
        # load db
        dbm = DBManager()
        matching_records = dbm.view_entries_for_duration(time_spent)
        if len(matching_records) == 0:
            print("\nNo matches, returning to search menu")
            return self.search_entries
        self.records = matching_records
        self.current_record = 0
        return self.present_next_result

    def search_text_search(self):
        """This is the menu where the user enters a text string and is presented
        with all entries containing that string in the task name or notes
        """
        print('SEARCH USING TEXT STRING')
        print("Enter the text string to search on")
        input_text = input("> ")
        text_string = input_text
        # load db
        dbm = DBManager()
        matching_records = dbm.view_entries_with_text(text_string)
        if len(matching_records) == 0:
            print("\nNo matches, returning to search menu")
            return self.search_entries
        self.records = matching_records
        self.current_record = 0
        return self.present_next_result

    # Modification Methods
    # --------------------
    def edit_record(self):
        print("edit record")
        print('enter the record number to edit')
        user_input = input("> ")
        match_index = int(user_input) - 1
        record = self.records[match_index]
        # get the new values for the record
        print("New Username")
        input_text = input("Enter the username > ")
        username = input_text
        date = None
        while date is None:
            print("New date of the Task")
            user_entry = self.date_entry()
            if user_entry[0] is not None:  # error
                print(user_entry[0])
                continue
            else:
                date = user_entry[1]
                date_string = self.date_to_string(date, target='file')
        print("New name of the Task")
        input_text = input("Enter the name of the task > ")
        task_name = input_text
        time_spent = None
        while time_spent is None:
            print("New time spent")
            input_text = input("Enter a whole number of minutes (rounded) ")
            try:
                time_spent = int(input_text)
            except ValueError:
                print("Invalid value")
                continue
        print("New notes")
        input_text = input("(Optional, leave blank for none) ")
        notes = input_text
        # load the db
        dbm = DBManager()
        # old_entry = dbm.view_entries
        new_values = {
            'name': username,
            'date': date_string,
            'task_name': task_name,
            'duration': time_spent,
            'notes': notes
        }
        dbm.edit_entry(record, new_values)
        return self.main_menu

    def edit_current_record(self):
        print("edit record")
        match_index = self.current_record
        record = self.records[match_index]
        # get the new values for the record
        print("New Username")
        input_text = input("Enter the username > ")
        username = input_text
        date = None
        while date is None:
            print("New date of the Task")
            user_entry = self.date_entry()
            if user_entry[0] is not None:  # error
                print(user_entry[0])
                continue
            else:
                date = user_entry[1]
        print("New name of the Task")
        input_text = input("Enter the name of the task > ")
        task_name = input_text
        time_spent = None
        while time_spent is None:
            print("New time spent")
            input_text = input("Enter a whole number of minutes (rounded) ")
            try:
                time_spent = int(input_text)
            except ValueError:
                print("Invalid value")
                continue
        print("New notes")
        input_text = input("(Optional, leave blank for none) ")
        notes = input_text
        # load the db
        dbm = DBManager()
        new_values = {
            'name': username,
            'date': date,
            'task_name': task_name,
            'duration': time_spent,
            'notes': notes
        }
        dbm.edit_entry(record, new_values)
        return self.main_menu

    def select_detail(self):
        print("View record")
        print('enter the record number to view')
        user_input = input("> ")
        match_index = int(user_input) - 1
        self.current_record = match_index
        return self.present_next_result

    def delete_record(self):
        print("delete record")
        print('enter the record number to delete')
        user_input = input("> ")
        match_index = int(user_input) - 1
        record = self.records[match_index]
        print(record)
        # load db
        dbm = DBManager()
        dbm.delete_entry(record)
        print("Entry deleted")
        return self.main_menu

    def delete_current_record(self):
        print("delete record")
        match_index = self.current_record
        record = self.records[match_index]
        # load db
        dbm = DBManager()
        dbm.delete_entry(record)
        print("Entry deleted")
        return self.main_menu

    # Other UI Methods
    # ----------------
    def display_entry(self, entry, verbose=False, return_only=False):
        """This method displays a selected entry, showing:
        - date (read from file in iso 8601 and displayed in whatever is set in
          options)
        - task name
        - time taken
        - any notes
        """
        username = entry[settings.HEADERS['user']]
        date_object = entry[settings.HEADERS['date']]
        date = self.date_to_string(date_object, target="display")
        task_name = entry[settings.HEADERS['task_name']]
        time_taken = entry[settings.HEADERS['duration']]
        notes = entry[settings.HEADERS['notes']]
        if verbose:
            line0 = username
            print(line0)
            line1 = "{}: {}".format(date, task_name)
            print(line1)
            print("-" * len(line1))
            print("{} minutes".format(time_taken))
            print("{}".format(notes))
        else:
            short_form = "{}: {} ({}m): {} | {}".format(username,
                                                        date,
                                                        time_taken,
                                                        task_name,
                                                        notes)
            if return_only:
                return short_form
            else:
                print(short_form)

    def previous_result(self):
        """load previous result"""
        self.current_record -= 1
        return self.present_next_result

    def next_result(self):
        """load next result"""
        self.current_record += 1
        return self.present_next_result

    def previous_page(self):
        """load previous page of results"""
        self.current_page_start -= self.OPTIONS['entries per page']
        return self.present_results

    def next_page(self):
        """load next page of results"""
        self.current_page_start += self.OPTIONS['entries per page']
        return self.present_results

    # Helper Methods
    # --------------
    def validate_date_entry(self, date_string, date_format):
        """Takes a date_string and date_format and attempts to create
        a valid datetime object with those imports.
        Returns a tuple in the form (error, datetime) where:
        - `error` is None if valid and a description of the error text if
          invalid;
        - `datetime` is a datetime object if valid and None if invalid
        """
        try:
            naive_datetime = datetime.datetime.strptime(
                date_string,
                date_format['datetime format']
            )
        except ValueError:
            error_text = "{date_string} is not valid in format {date_format}"
            error_args = {"date_string": date_string,
                          "date_format": date_format['UI format']}
            return (error_text.format(**error_args), None)
        else:
            if not self.OPTIONS['allow future dates']:
                if naive_datetime > datetime.datetime.now():
                    error_text = "dates in the future are not permitted"
                    error_args = {"date_string": date_string,
                                  "date_format": date_format['UI format']}
                    return (error_text.format(**error_args), None)
                if naive_datetime < self.OPTIONS['earliest allowed date']:
                    bad_date = self.OPTIONS['earliest allowed date'].strftime(
                        self.OPTIONS['date format']['datetime format']
                    )
                    error_text = "dates before {} are not permitted".format(
                        bad_date
                    )
                    error_args = {"date_string": date_string,
                                  "date_format": date_format['UI format']}
                    return (error_text.format(**error_args), None)
            return (None, naive_datetime)

    def date_entry(self):
        """This helper function asks for a date input in the user's preferred
        format and then returns that date as a naive datetime object
        """
        date_format = self.OPTIONS['date format']
        input_text = "Please use the '{}' date format: "
        user_entry = input(input_text.format(date_format['UI format']))
        # validate date entry
        validated = self.validate_date_entry(user_entry, date_format)
        return validated

    def date_to_string(self, date_object, target='display'):
        """This helper function takes a naive date object and returns a
        string representation in:
        - `target='display'`: the user's preferred display format
        - `target='file': the save format
        """
        if target == 'display':
            option = self.OPTIONS['date format']
            string_format = option['datetime format']
        else:  # 'file' or unrecognised target, fallback to write mode
            option = self.OPTIONS['save format (date)']
            string_format = option['datetime format']
        return date_object.strftime(string_format)

# ---------------------------

if __name__ == "__main__":

    menu = Menu()


import math
import torch
import gpytorch
import numpy as np
from voltron.means import EWMAMean, DEWMAMean, TEWMAMean
from botorch.models import SingleTaskGP
from botorch.optim.fit import fit_gpytorch_torch
from voltron.rollout_utils import nonvol_rollouts

class BasicGP():
    def __init__(self, train_x, train_y, kernel="matern", mean='constant',
                 k=400, num_mixtures=10):
#         super(BasicGP, self).__init__(train_x, train_y, likelihood)
        if mean.lower() == 'constant':
            mean_module = gpytorch.means.ConstantMean().to(train_x.device)
        elif mean.lower() == 'ewma':
            mean_module = EWMAMean(train_x, train_y, k).to(train_x.device)
        elif mean.lower() == 'dewma':
            mean_module = DEWMAMean(train_x, train_y, k).to(train_x.device)
        elif mean.lower() == 'tewma':
            mean_module = TEWMAMean(train_x, train_y, k).to(train_x.device)
        else:
            print("ERROR: Mean not implemented")

        if kernel.lower() == 'matern':
            covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.MaternKernel())
        elif kernel.lower() in ['sm', 'spectralmixture', 'spectral']:
            covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=num_mixtures)
            covar_module.initialize_from_data(train_x, train_y)
        elif kernel.lower() == 'rbf':
            covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
        else:
            print("ERROR: Kernel not implemented")

        self.model = SingleTaskGP(train_x.view(-1, 1), train_y.reshape(-1, 1),
                                  covar_module=covar_module,
                                  likelihood=gpytorch.likelihoods.GaussianLikelihood())
        self.model.mean_module = mean_module

    def Train(self, train_iters=400, display=False):
        mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.model.likelihood, self.model)
        fit_gpytorch_torch(mll, options={'maxiter':train_iters, 'disp':display})

    def Forecast(self, test_x, nsample=100):
        if not isinstance(self.model.mean_module, (EWMAMean, DEWMAMean, TEWMAMean)):

            samples = self.model.posterior(test_x).sample(torch.Size((nsample, )))

        else:
            samples = nonvol_rollouts(self.model.train_inputs[0].squeeze(),
                                      self.model.train_targets.squeeze(),
                                      test_x, self.model, nsample)
        return samples.squeeze()


import gspread
import pandas as pd
import os
import requests
from google.oauth2 import service_account
# from oauth2client.service_account import ServiceAccountCredentials
from bs4 import BeautifulSoup

scope = ['https://spreadsheets.google.com/feeds',
         'https://www.googleapis.com/auth/drive']
secret_file = os.path.join(os.getcwd(), 'client_secret.json')

credentials = service_account.Credentials.from_service_account_file(
    secret_file, scopes=scope)

gc = gspread.authorize(credentials)

sht1 = gc.open_by_key('1Jgsf-5wtsCdyDiIp-P_EZbnoT6Ya_0URC3I8s1-5GeM')

df = pd.DataFrame(sht1.worksheet("Fontes").get_all_values()[1:])
df.columns = df.iloc[0]
df.drop(df.index[0], inplace=True)

header = {"User-Agent": "Mozilla/5.0 (Windows NT 6.1; Win64; x64)"
          "AppleWebKit/537.36 (KHTML, like Gecko)"
          "Chrome/54.0.2840.71 Safari/537.36",
          "upgrade-insecure-requests": "1",
          "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,"
          "image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3",
          "accept-encoding": "gzip, deflate, br",
          "accept-language": "en-GB,en-US;q=0.9,en;q=0.8",
          "cache-control": "max-age=0"}

res = requests.get(df['Competitor Link'].iloc[0], headers=header)

soup = BeautifulSoup(res.text, "html.parser")

price = soup.find("div", {"class": "preco_normal"}).text.replace("R$", "")
print(price)
sht1.worksheet("Fontes").update_acell('C3', price)


# -*- coding: utf-8 -*-
#
# Copyright (c) 2012 David Townshend
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
# for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 675 Mass Ave, Cambridge, MA 02139, USA.

from norman._six import assert_raises

from norman import NotSet, Field, Store, Index

try:
    from unittest.mock import Mock, patch
except ImportError:
    from mock import Mock, patch

##### Index tests

class TestIndex_Ordered(object):
    'Tests starting with an empty index'

    def setup(self):
        field = Mock(key=lambda x: x)
        self.i = Index(field)
        self.orecords = ['M' + str(i) for i in range(6)]
        self.urecords = ['U0', 'U1', 'U2']
        self.i._ordered = ([0, 1, 2, 3, 3, 4], [r for r in self.orecords])
        self.ordered = ([0, 1, 2, 3, 3, 4], [r for r in self.orecords])
        self.unordered = self.i._unordered.copy()
        self.i._unordered[-1] = [('1', self.urecords[0]),
                                 ('1', self.urecords[1])]
        self.i._unordered[-2] = [('2', self.urecords[2])]
        self.unordered[-1] = [('1', self.urecords[0]),
                              ('1', self.urecords[1])]
        self.unordered[-2] = [('2', self.urecords[2])]

    def test_insert(self):
        r = Mock()
        self.i.insert(1, r)
        self.orecords.insert(2, r)
        assert self.i._ordered == ([0, 1, 1, 2, 3, 3, 4], self.orecords)
        assert self.i._unordered == self.unordered

    def test_len(self):
        assert len(self.i) == 9, len(self.i)

    def test_remove(self):
        self.i.remove(2, self.orecords[2])
        del self.orecords[2]
        expect = ([0, 1, 3, 3, 4], self.orecords)
        assert self.i._ordered == expect
        assert self.i._unordered == self.unordered

    def test_iter_eq(self):
        got = set(self.i == 3)
        expect = set(self.orecords[3:5])
        assert got == expect, (got, expect)

    def test_iter_ne(self):
        expect = set(self.orecords[:3] + self.orecords[5:] + self.urecords)
        got = set(self.i != 3)
        assert got == expect, (got, expect)

    def test_iter_lt(self):
        got = set(self.i < 3)
        expect = set(self.orecords[:3])
        assert got == expect, (got, expect)

    def test_iter_le(self):
        got = set(self.i <= 3)
        expect = set(self.orecords[:5])
        assert got == expect, (got, expect)

    def test_iter_gt(self):
        got = set(self.i > 2)
        expect = set(self.orecords[3:])
        assert got == expect, (got, expect)

    def test_iter_ge(self):
        got = set(self.i >= 2)
        expect = set(self.orecords[2:])
        assert got == expect, (got, expect)

class TestIndex_UnOrdered(object):
    'Tests starting with an empty index'

    def setup(self):
        def key(value):
            raise TypeError
        field = Mock(key=key)
        self.i = Index(field)
        self.orecords = ['M' + str(i) for i in range(6)]
        self.urecords = ['U0', 'U1', 'U2']
        self.i._ordered = ([0, 1, 2, 3, 3, 4], [r for r in self.orecords])
        self.ordered = ([0, 1, 2, 3, 3, 4], [r for r in self.orecords])
        self.unordered = self.i._unordered.copy()
        self.i._unordered[-1] = [('1', self.urecords[0]),
                                 ('1', self.urecords[1])]
        self.i._unordered[-2] = [('2', self.urecords[2])]
        self.unordered[-1] = [('1', self.urecords[0]),
                              ('1', self.urecords[1])]
        self.unordered[-2] = [('2', self.urecords[2])]

        def mockhash(v):
            try:
                return -int(v)
            except ValueError:
                raise TypeError
        patch('norman._store.hash', mockhash, create=True).start()
        patch('norman._store.id', lambda v: 10, create=True).start()

    def teardown(self):
        patch.stopall()

    def test_insert1(self):
        r = Mock()
        self.i.insert('4', r)
        self.unordered[-4] = [('4', r)]
        assert self.i._unordered == self.unordered
        assert self.i._ordered == self.ordered

    def test_insert2(self):
        r = Mock()
        self.i.insert('2', r)
        self.unordered[-2].append(('2', r))
        assert self.i._unordered == self.unordered
        assert self.i._ordered == self.ordered

    def test_insert_id(self):
        r = Mock()
        self.i.insert('a', r)
        self.unordered[10] = [('a', r)]
        assert self.i._unordered == self.unordered
        assert self.i._ordered == self.ordered

    def test_insert_NotSet(self):
        r = Mock()
        self.i.insert(NotSet, r)
        self.unordered[NotSet] = [(NotSet, r)]
        assert self.i._unordered == self.unordered
        assert self.i._ordered == self.ordered

    def test_remove(self):
        self.i.remove('1', self.urecords[1])
        self.unordered[-1].pop()
        assert self.i._unordered == self.unordered
        assert self.i._ordered == self.ordered

    def test_remove_id(self):
        self.i.remove('2', self.urecords[2])
        del self.unordered[-2]
        assert self.i._unordered == self.unordered
        assert self.i._ordered == self.ordered

    def test_iter_eq(self):
        expect = set(self.urecords[:2])
        got = set(self.i == '1')
        assert got == expect, (got, expect)

    def test_iter_ne(self):
        expect = set([self.urecords[2]] + self.orecords)
        got = set(self.i != '1')
        assert got == expect, (got, expect)

    def test_comparison(self):
        with assert_raises(TypeError):
            set(self.i < '1')
        with assert_raises(TypeError):
            set(self.i <= '1')
        with assert_raises(TypeError):
            set(self.i > '1')
        with assert_raises(TypeError):
            set(self.i >= '1')

class TestIndexCornerCases(object):

    def test_notset(self):
        field = Field()
        i = Index(field)
        r = Mock()
        i.insert(NotSet, r)
        assert i._unordered[NotSet] == [(NotSet, r)]

class TestStore(object):

    def setup(self):
        self.full = Field(default=NotSet)
        self.sparse = Field(default= -1)
        self.missing = Field(default=NotSet)
        self.full._name = 'full'
        self.sparse._name = 'sparse'
        self.missing._name = 'missing'
        self.store = Store()
        self.store.add_field(self.full)
        self.store.add_field(self.sparse)

    def populate(self):
        for i in range(5):
            self.store.add_record(str(i))
            self.store.set(str(i), self.full, i)
        for i in [1, 3]:
            self.store.set(str(i), self.sparse, i)

    def test_add_field(self):
        self.populate()
        self.store.add_field(self.missing)
        assert self.store.get('0', self.missing) == NotSet

    def test_add_record(self):
        self.store.add_record('new')
        assert self.store.has_record('new')

    def test_clear(self):
        self.populate()
        self.store.clear()
        assert self.store.record_count() == 0
        for index in self.store.indexes.values():
            assert len(index) == 0

    def test_get(self):
        self.populate()
        assert self.store.get('0', self.full) == 0

    def test_has_record(self):
        self.populate()
        assert self.store.has_record('0')
        assert not self.store.has_record('not a record')

    def test_iter_field_full(self):
        self.populate()
        got = set(self.store.iter_field(self.full))
        expect = set([(str(i), i) for i in range(5)])
        assert got == expect, got

    def test_iter_field_sparse(self):
        self.populate()
        got = set(self.store.iter_field(self.sparse))
        expect = set([('0', -1), ('1', 1), ('2', -1), ('3', 3), ('4', -1)])
        assert got == expect, got

    def test_iter_records(self):
        self.populate()
        it = self.store.iter_records()
        assert set(it) == set('012324')

    def test_record_count(self):
        self.populate()
        assert self.store.record_count() == 5

    def test_remove_record(self):
        self.populate()
        self.store.remove_record('3')
        assert set(self.store.iter_records()) == set('0124')

    def test_remove_field(self):
        # This merely tests that it runs.
        self.populate()
        self.store.remove_field(self.full)

    def test_set_overwrite(self):
        self.populate()
        self.store.set('1', self.sparse, 'new value')
        assert self.store.get('1', self.sparse) == 'new value'

class TestStoreIndex(object):

    def setup(self):
        self.f = Field()
        self.f._name = 'f'
        self.store = Store()
        self.store.add_field(self.f)
        self.index = self.store.indexes[self.f]

    def test_add_record(self):
        self.store.add_record('new')
        assert self.index._unordered[NotSet] == [(NotSet, 'new')]
        assert self.index._ordered == ([], [])

    def test_remove_record(self):
        self.store.add_record('new')
        self.store.remove_record('new')
        assert self.index._unordered == {}
        assert self.index._ordered == ([], [])

    def test_set(self):
        self.store.add_record('new')
        self.store.set('new', self.f, 'value')
        assert self.index._unordered == {}
        assert self.index._ordered == ([('1str', 'value')], ['new'])


import pygame, sys, time, random
from pygame.locals import *

# Установка pygame.
pygame.init()
mainClock = pygame.time.Clock()

# Настройка окна.
WINDOWWIDTH = 400
WINDIWHEIGHT = 400
windowSurface = pygame.display.set_mode((WINDOWWIDTH, WINDIWHEIGHT), 0, 32)
pygame.display.set_caption('Спрайты и звуки')

# Настройка цвета.
WHITE = (255, 255, 255)

# Создание структуры данных блока.
player = pygame.Rect(300, 100, 40, 40)
playerImage = pygame.image.load('player.png')
playerStretchedImage = pygame.transform.scale(playerImage, (40, 40))
foodImage = pygame.image.load('cherry.png')
foods = []
for i in range(20):
    foods.append(pygame.Rect(random.randint(0, WINDOWWIDTH - 20),
                             random.randint(0, WINDIWHEIGHT - 20), 20, 20))

foodCounter = 0
NEWFOOD = 40

# Создание переменных клавиатуры.
moveLeft = False
moveRight = False
moveUp = False
moveDown = False

MOVESPEED = 6

# Настройка музыка.
pickUpSound = pygame.mixer.Sound('pickup.wav')
pygame.mixer.music.load('gameover.wav')
pygame.mixer.music.play(-1, 0.0)
musicPlaying = True

# Запуск игрового цикла.
while True:
    # Проверка наличия события QUIT.
    for event in pygame.event.get():
        if event.type == QUIT:
            pygame.quit()
            sys.exit()

        if event.type == KEYDOWN:
            # Изменение переменных клавиатуры.
            if event.key == K_LEFT or event.key == K_a:
                moveRight = False
                moveLeft = True
            if event.key == K_RIGHT or event.key == K_d:
                moveLeft = False
                moveRight = True
            if event.key == K_UP or event.key == K_w:
                moveDown = False
                moveUp = True
            if event.key == K_DOWN or event.key == K_s:
                moveUp = False
                moveDown = True
        if event.type == KEYUP:
            if event.key == K_ESCAPE:
                pygame.quit()
                sys.exit()
            if event.key == K_LEFT or event.key == K_a:
                moveLeft = False
            if event.key == K_RIGHT or event.key == K_d:
                moveRight = False
            if event.key == K_UP or event.key == K_w:
                moveUp = False
            if event.key == K_DOWN or event.key == K_s:
                moveDown = False
            if event.key == K_x:
                player.top = random.randint(0, WINDIWHEIGHT - player.height)
                player.left = random.randint(0, WINDOWWIDTH - player.width)
            if event.key == K_m:
                if musicPlaying:
                    pygame.mixer.music.stop()
                else:
                    pygame.mixer.music.play(-1, 0.0)
                musicPlaying = not musicPlaying

        if event.type == MOUSEBUTTONUP:
            foods.append(pygame.Rect(event.pos[0] - 10, event.pos[1] - 10,
                                     20, 20))

    foodCounter += 1
    if foodCounter >= NEWFOOD:
        # Добавление новой "еды".
        foodCounter = 0
        foods.append(pygame.Rect(random.randint(0, WINDOWWIDTH - 20),
                                 random.randint(0, WINDIWHEIGHT - 20),
                                 20, 20))

        # Создание на поверхности белого фона.
        windowSurface.fill(WHITE)

        # Перемещение игрока.
        if moveDown and player.bottom < WINDIWHEIGHT:
            player.top += MOVESPEED
        if moveUp and player.top > 0:
            player.top -= MOVESPEED
        if moveLeft and player.left > 0:
            player.left -= MOVESPEED
        if moveRight and player.right < WINDOWWIDTH:
            player.right += MOVESPEED

        # Отображение блока на поверхности.
        windowSurface.blit(playerStretchedImage, player)

        # Проверка, не пересекся ли блок с каким-либо блоками "еды".
        for food in foods[:]:
            if player.colliderect(food):
                foods.remove(food)
                player = pygame.Rect(player.left, player.top,
                                     player.width + 2, player.height + 2)
                playerStretchedImage = pygame.transform.scale(playerImage,
                                                             (player.width,
                                                              player.height))
            if musicPlaying:
                pickUpSound.play()

        # Отображение "еды".
        for food in foods:
            windowSurface.blit(foodImage, food)

        # Вывод окна на экран.
        pygame.display.update()
        mainClock.tick(40)


from django.conf.urls.static import static
from django.contrib.auth.decorators import login_required
from django.urls import path
from . import views
from .views import AddPostView, UserSettings, UserProfile
from .models import Profile
from newsletter.models import NewsLetter

urlpatterns = [
    path(
        "",
        views.PostList.as_view(),
        name="home"
        ),
    path(
        'featured/',
        views.FeaturedView.as_view(),
        name='featured'
        ),
    path(
        'featured_post/',
        views.FeaturedPost,
        name='featured_post'
        ),
    path(
        'profile/',
        login_required(UserProfile.as_view()),
        name='profile'
        ),
    path(
        '<pk>/delete_profile/',
        login_required(views.DeleteProfile.as_view()),
        name='delete_profile'
        ),
    path(
        'edit_profile/',
        login_required(UserSettings.as_view()),
        name='edit_profile'
        ),
    path(
        'add_post/',
        AddPostView.as_view(),
        name='add_post'
        ),
    path(
        'pets/',
        views.PetsPost,
        name='pets'
        ),
    path(
        '<slug:slug>/',
        views.PostDetail.as_view(),
        name='post_detail'
        ),
    path(
        'like/<slug:slug>/',
        views.PostLike.as_view(),
        name='post_like'
        ),
    path(
        '<int:id>/delete-post',
        views.delete_post,
        name='delete_post'
        ),
    path(
        'edit/<int:post_id>',
        views.edit_post,
        name='edit_post'
        ),
]


from collections import deque

def bfs(graph, root):
    visited = set()
    queue = deque([root])
    i = 0
    while queue:
        n = queue.popleft()
        if n not in visited:
            if i != 0:
                visited.add(n)
            queue += set(graph[n]) - set(visited)
        i += 1

    return visited

n = int(input())
lst = []
dic = {}

for _ in range(n):
    lst.append(list(map(int, input().split())))

for i in range(n):
    dic[i] = []
    for j in range(n):
        if lst[i][j] == 1:
            a = dic[i]
            a.append(j)

lst_ans = [[0]*n for _ in range(n)]
for i in range(n):
    visited_current = bfs(dic, i)
    for node in visited_current:
        lst_ans[i][node] = 1

for i in range(n):
    for j in range(n):
        print(lst_ans[i][j], end=" ")
    print()

# 플로이드 - 와샬을 쓰면 더 쉽다.
# 플로이드 - 와샬 쓰이는 데가 많아 보인다.