#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
programs.py
----------------------------------
the different source programs
'''

import csv
import re
from . import schema
import sqlite3
import os
from collections import OrderedDict
from datetime import datetime
from dateutil.parser import parse as dateparser
from glob import glob
from os.path import join, isdir, basename, splitext
from .querycsv import query_csv
from functools import partial
from .services import Caster, Normalizer, ChargeBalancer, HttpClient
from .benchmarking import get_milliseconds

class Program(object):
    most_recent_result_query = ('SELECT max(SampleDate) FROM [UGSWaterChemistry].[ugswaterchemistry].[Results]'
                                ' WHERE [DataSource] = \'{}\'')
    new_stations_query = ('SELECT * FROM (VALUES{}) AS t(StationId) WHERE NOT EXISTS('
                          'SELECT 1 FROM [UGSWaterChemistry].[ugswaterchemistry].[Stations] WHERE [StationId] = t.StationId)')

    def _get_most_recent_result_date(self, datasource):

        #: open connection if one hasn't been opened
        if not hasattr(self, 'cursor') or not self.cursor:
            self.cursor = self.cursor_factory(self.db['connection_string'])

        query = self.most_recent_result_query.format(datasource)
        try:
            last_updated = self.cursor.execute(query).fetchone()
        except Exception as e:
            del self.cursor
            raise e

        #: fetchone returns a set with one item
        if last_updated and len(last_updated) == 1:
            return last_updated[0]

    def _get_unique_station_ids(self, station_ids):
        '''queries the Stations table to find stations that have not been inserted yet
        stations_ids: list('StationId')

        returns a set of station ids
        '''

        station_ids = ['(\'{}\')'.format(station_id) for station_id in station_ids]

        if not hasattr(self, 'cursor') or not self.cursor:
            self.cursor = self.cursor_factory(self.db['connection_string'])

        statement = self.new_stations_query.format(','.join(station_ids))
        self.cursor.execute(statement)

        return self.cursor.fetchall()

class WqpProgram(Program):
    '''class for handling wqp csv files'''

    datasource = 'WQP'

    TEMPDB = 'temp.sqlite3'

    wqp_url = ('http://www.waterqualitydata.us/{}/search?sampleMedia=Water&startDateLo={}&startDateHi={}&'
               'bBox=-115%2C35.5%2C-108%2C42.5&mimeType=csv')

    fields = {'sample_id': 'ActivityIdentifier', 'monitoring_location_id': 'MonitoringLocationIdentifier'}

    sql = {
        'distinct_sample_id': 'select distinct({}) from {}',
        'sample_id': 'select * from {} where {} = \'{}\'',
        'wqxids': 'select {0} from {1} where {0} LIKE \'%_WQX%\'',
        'new_results': ('SELECT * FROM (VALUES{}) AS t(SampleId) WHERE NOT EXISTS('
                        'SELECT 1 FROM [UGSWaterChemistry].[ugswaterchemistry].[Results] WHERE [SampleId] = t.SampleId)'),
        'create_index': 'CREATE INDEX IF NOT EXISTS "ActivityIdentifier_{0}" ON "{0}" ("ActivityIdentifier" ASC)'
    }

    wqx_re = re.compile('(_WQX)-')

    station_config = OrderedDict([
        ('OrganizationIdentifier', 'OrgId'),
        ('OrganizationFormalName', 'OrgName'),
        ('MonitoringLocationIdentifier', 'StationId'),
        ('MonitoringLocationName', 'StationName'),
        ('MonitoringLocationTypeName', 'StationType'),
        ('MonitoringLocationDescriptionText', 'StationComment'),
        ('HUCEightDigitCode', 'HUC8'),
        ('LongitudeMeasure', 'Lon_X'),
        ('LatitudeMeasure', 'Lat_Y'),
        ('HorizontalAccuracyMeasure/MeasureValue', 'HorAcc'),
        ('HorizontalAccuracyMeasure/MeasureUnitCode', 'HorAccUnit'),
        ('HorizontalCollectionMethodName', 'HorCollMeth'),
        ('HorizontalCoordinateReferenceSystemDatumName', 'HorRef'),
        ('VerticalMeasure/MeasureValue', 'Elev'),
        ('VerticalMeasure/MeasureUnitCode', 'ElevUnit'),
        ('VerticalAccuracyMeasure/MeasureValue', 'ElevAcc'),
        ('VerticalAccuracyMeasure/MeasureUnitCode', 'ElevAccUnit'),
        ('VerticalCollectionMethodName', 'ElevMeth'),
        ('VerticalCoordinateReferenceSystemDatumName', 'ElevRef'),
        ('StateCode', 'StateCode'),
        ('CountyCode', 'CountyCode'),
        ('AquiferName', 'Aquifer'),
        ('FormationTypeText', 'FmType'),
        ('AquiferTypeName', 'AquiferType'),
        ('ConstructionDateText', 'ConstDate'),
        ('WellDepthMeasure/MeasureValue', 'Depth'),
        ('WellDepthMeasure/MeasureUnitCode', 'DepthUnit'),
        ('WellHoleDepthMeasure/MeasureValue', 'HoleDepth'),
        ('WellHoleDepthMeasure/MeasureUnitCode', 'HoleDUnit'),
        ('demELEVm', '!demELEVm'),
        ('DataSource', '!DataSource'),
        ('WIN', '!WIN'),
        ('Shape', '!Shape')
    ])

    result_config = OrderedDict([
        ('AnalysisStartDate', 'AnalysisDate'),
        ('ResultAnalyticalMethod/MethodName', 'AnalytMeth'),
        ('ResultAnalyticalMethod/MethodIdentifier', 'AnalytMethId'),
        ('AutoQual', '!AutoQual'),
        ('CAS_Reg', '!CAS_Reg'),
        ('Chrg', '!Chrg'),
        ('DataSource', '!DataSource'),
        ('ResultDetectionConditionText', 'DetectCond'),
        ('IdNum', '!IdNum'),
        ('ResultLaboratoryCommentText', 'LabComments'),
        ('LaboratoryName', 'LabName'),
        ('Lat_Y', '!Lat_Y'),
        ('DetectionQuantitationLimitTypeName', 'LimitType'),
        ('Lon_X', '!Lon_X'),
        ('DetectionQuantitationLimitMeasure/MeasureValue', 'MDL'),
        ('DetectionQuantitationLimitMeasure/MeasureUnitCode', 'MDLUnit'),
        ('MethodDescriptionText', 'MethodDescript'),
        ('OrganizationIdentifier', 'OrgId'),
        ('OrganizationFormalName', 'OrgName'),
        ('CharacteristicName', 'Param'),
        ('ParamGroup', '!ParamGroup'),
        ('ProjectIdentifier', 'ProjectId'),
        ('MeasureQualifierCode', 'QualCode'),
        ('ResultCommentText', 'ResultComment'),
        ('ResultStatusIdentifier', 'ResultStatus'),
        ('ResultMeasureValue', 'ResultValue'),
        ('ActivityCommentText', 'SampComment'),
        ('ActivityDepthHeightMeasure/MeasureValue', 'SampDepth'),
        ('ActivityDepthAltitudeReferencePointText', 'SampDepthRef'),
        ('ActivityDepthHeightMeasure/MeasureUnitCode', 'SampDepthU'),
        ('SampleCollectionEquipmentName', 'SampEquip'),
        ('ResultSampleFractionText', 'SampFrac'),
        ('ActivityStartDate', 'SampleDate'),
        ('ActivityStartTime/Time', 'SampleTime'),
        ('ActivityIdentifier', 'SampleId'),
        ('ActivityMediaSubdivisionName', 'SampMedia'),
        ('SampleCollectionMethod/MethodIdentifier', 'SampMeth'),
        ('SampleCollectionMethod/MethodName', 'SampMethName'),
        ('ActivityTypeCode', 'SampType'),
        ('MonitoringLocationIdentifier', 'StationId'),
        ('ResultMeasure/MeasureUnitCode', 'Unit'),
        ('USGSPCode', 'USGSPCode')
    ])

    def __init__(self,
                 logger,
                 db,
                 update,
                 source=None,
                 secrets=None,
                 sql_statements={},
                 update_row=None,
                 insert_rows=None,
                 cursor_factory=None,
                 arcpy=None):
        self.logger = logger
        '''create a new WQP program
        db - the connection string for the database to seed
        source - the path on disk to find csv files to ETL
        update - boolean value whether we are seeding or updating
        secrets - ignored. Valid for sdwis only.
        sql_statements - common sql statements for inserting into stations and results
        update_row - the function to reproject a point and set the DataSource
        insert_rows - the function to batch insert rows
        cursor_factory - the function to create pyodbc connection_string
        arcpy - ignored. for gdb programs only
        '''
        self.db = db
        self._update_row = update_row
        self._insert_rows = insert_rows
        self.cursor_factory = cursor_factory
        self.sql.update(sql_statements)

        #: if update is True then we are updating
        if not update:
            #: check that source exists wqp/results and wqp/stations
            parent_folder = join(source, self.datasource)

            if not isdir(parent_folder):
                raise Exception('Pass in a location to the parent folder that contains WQP. {}'.format(parent_folder))

            results_folder = join(parent_folder, 'Results')
            stations_folder = join(parent_folder, 'Stations')

            if not isdir(results_folder or stations_folder):
                raise Exception('WQP folder is missing Results or Stations child folders containing csv files.')

            #: all is well set the folders for seeding later on
            self.results_folder = results_folder
            self.stations_folder = stations_folder

    def seed(self):
        if not hasattr(self, 'results_folder') or not hasattr(self, 'stations_folder'):
            raise Exception('You must pass a file location if you are seeding. Did you want to update?')

        try:
            self._seed_by_file()
        finally:
            if hasattr(self, 'cursor'):
                del self.cursor

    def update(self):
        self.logger.info('updating {} stations and results...'.format(self.datasource))

        try:
            last_updated = self._get_most_recent_result_date(self.datasource)

            if not last_updated:
                raise Exception('No last updated date. You should seed some data first.')

            self.logger.info('fetching records after {}'.format(last_updated))
            #: get new results from wqp service
            result_url = self._format_url(self.wqp_url, 'Result', last_updated)
            #: group them as if they were read from querycsv
            new_results = self._group_rows_by_id(HttpClient.get_csv(result_url, self.logger))
            #: remove results that have a sample id already in the database
            new_results = self._remove_existing_results(new_results)
            #: find the station ids from the new results that aren't in the database
            new_station_ids = self._find_new_station_ids(new_results)
            #: check database for stripped wqx and remove wqx id's since they are duplicates
            new_station_ids = self._remove_existing_wqx_station_ids(new_station_ids)

            if new_station_ids and len(new_station_ids) > 0:
                self.logger.info('of the new stations found, attempting to insert {}'.format(len(new_station_ids)))

                station_url = self._format_url(self.wqp_url, 'Station', last_updated)
                new_stations = HttpClient.get_csv(station_url, self.logger)

                if not new_stations:
                    raise Exception('WQP service should have returned results but result is empty. {}'.format(
                        station_url))

                header = next(new_stations)

                stations = self._extract_stations_by_id(new_stations, new_station_ids, header)
                wqx = self._get_wqx_duplicate_ids(stations)

                self._seed_stations(stations, header=header, wqx=wqx)
            else:
                self.logger.info('all stations already in database')
            for samples_for_id in list(new_results.values()):
                self._seed_results(samples_for_id)

        finally:
            if hasattr(self, 'cursor'):
                del self.cursor

    def _seed_by_file(self):
        self.logger.info('processing {} stations...'.format(self.datasource))

        for csv_file in self._get_files(self.stations_folder):
            #: create csv reader
            with open(csv_file, 'r') as f:
                self.logger.info('- {}'.format(basename(csv_file)))

                #: generate duplicate id list
                wqx = self._get_wqx_duplicate_ids(csv_file)

                reader = csv.reader(f)
                #: get a reference to the header row
                header = next(reader)

                self._seed_stations(reader, header=header, wqx=wqx)

                self.logger.info('- {}: done'.format(basename(csv_file)))

        self.logger.info('processing {} stations done.'.format(self.datasource))
        self.logger.info('processing {} results...'.format(self.datasource))

        for csv_file in self._get_files(self.results_folder):
            self.logger.info('- {}'.format(basename(csv_file)))

            #: create sqlite db and get unique sample ids
            unique_sample_ids = self._get_distinct_sample_ids_from(csv_file)

            try:
                #: this needs to be done after get_distinct so that the table is created in the db
                self._add_sample_index(csv_file)

                sample_sets = 0
                sets_start = get_milliseconds()

                for sample_id in unique_sample_ids:
                    samples = self._get_samples_for_id(sample_id, csv_file)

                    self._seed_results(samples)

                    sample_sets += 1
                    elapsed = get_milliseconds() - sets_start
                    if sample_sets % 5000 == 0:
                        self.logger.info('{} total sample sets (avg {} milliseconds per set)'.format(sample_sets, round(
                            elapsed / sample_sets, 5)))

            finally:
                #: in case something goes wrong always clean up the db
                os.remove(self.TEMPDB)
            self.logger.info('- {}: done'.format(basename(csv_file)))

        self.logger.info('processing {} results done.'.format(self.datasource))

    def _seed_stations(self, rows, header=None, wqx=None):
        stations = []

        for row in rows:
            #: push all the csv column names into the standard names
            row = self._etl_column_names(row, self.station_config, header=header)

            #: check for duplicate stationid and skip if found
            station_id = row['StationId']
            if not self.wqx_re.search(station_id) and station_id in wqx:
                continue

            #: cast
            row = Caster.cast(row, schema.station)

            #: set datasource, reproject and update shape
            row = self._update_row(row, self.datasource)

            #: normalize data including stripping _WXP etc
            row = Normalizer.normalize_station(row)

            #: reorder and filter out any fields not in the schema
            row = Normalizer.reorder_filter(row, schema.station)

            #: have to generate sql manually because of quoting on spatial WKT
            row = Caster.cast_for_sql(row)

            #: store row for later
            stations.append(list(row.values()))

        if not hasattr(self, 'cursor') or not self.cursor:
            self.cursor = self.cursor_factory(self.db['connection_string'])

        #: insert stations
        self._insert_rows(stations, self.sql['station_insert'], self.cursor)

    def _seed_results(self, samples_for_id):
        #: cast to defined schema types
        samples = list(map(partial(Caster.cast, schema=schema.result), samples_for_id))

        #: set datasource and spatial information
        samples = [self._update_row(sample, self.datasource) for sample in samples]

        #: normalize chemical names and units
        samples = list(map(Normalizer.normalize_sample, samples))

        #: create charge balance rows from sample
        charge_balances = ChargeBalancer.get_charge_balance(samples)

        samples.extend(charge_balances)

        #: reorder and filter out any fields not in the schema
        samples = list(map(partial(Normalizer.reorder_filter, schema=schema.result), samples))

        samples = list(map(Caster.cast_for_sql, samples))

        rows = [list(sample.values()) for sample in samples]

        if not hasattr(self, 'cursor') or not self.cursor:
            self.cursor = self.cursor_factory(self.db['connection_string'])

        #: TODO determine if this should this be batched in sets bigger than just a sample set?
        self._insert_rows(rows, self.sql['result_insert'], self.cursor)

    def _get_files(self, location):
        '''Takes the file location and returns the csv's within it.'''

        if not location:
            raise Exception('Pass in a location containing csv files to import.')

        files = glob(join(location, '*.csv'))

        if len(files) < 1:
            raise Exception(location, 'No csv files found.')

        return files

    def _get_distinct_sample_ids_from(self, file_path):
        '''Given a file_path, this returns a set of unique sample ids.'''

        file_name = self._get_file_name_without_extension(file_path)

        unique_sample_ids = query_csv(self.sql['distinct_sample_id'].format(self.fields['sample_id'], file_name),
                                      [file_path], self.TEMPDB)
        if len(unique_sample_ids) > 0:
            #: remove header cell
            unique_sample_ids.pop(0)

        return unique_sample_ids

    def _get_wqx_duplicate_ids(self, file_path):
        '''Given the file_path, return the list of stripped station ids that have the _WQX suffix'''

        file_name = None
        if isinstance(file_path, str):
            file_name = self._get_file_name_without_extension(file_path)
        else:
            #: we are sending in the updated stations
            stations = file_path

        if file_name:
            rows = query_csv(self.sql['wqxids'].format(self.fields['monitoring_location_id'], file_name), [file_path],
                             self.TEMPDB)
            if len(rows) > 0:
                rows.pop(0)

            return set([re.sub(self.wqx_re, '-', row[0]) for row in rows])

        return set([re.sub(self.wqx_re, '-', station['StationId'])
                    for station in [x for x in stations if self.wqx_re.search(x['StationId'])]])

    def _get_samples_for_id(self, sample_id_set, file_path, config=None):
        '''Given a `(id,)` styled set of sample ids, this will return the sample
        rows for the given csv file_path. The format will be an array of dictionaries
        with the key being the destination field name and the value being the source csv value.

        This is only invoked for result files.
        '''

        file_name = self._get_file_name_without_extension(file_path)
        samples_for_id = query_csv(self.sql['sample_id'].format(file_name, self.fields['sample_id'], sample_id_set[0]),
                                   [file_path], self.TEMPDB)

        return self._etl_column_names(samples_for_id, config or self.result_config)

    def _etl_column_names(self, rows, config, header=None):
        '''Given a dictionary or list of dictionaries, return a new row or
        list of rows with the correct field names'''

        #: we have an already etl'd station. skip it.
        if type(rows) is dict:
            return rows

        if len(rows) == 0:
            return None

        if not header:
            #: get header cell from rows and remove
            header = rows.pop(0)

        def return_value_if_not_in_config(key):
            if key in config:
                return config[key]

            return key

        header = [return_value_if_not_in_config(x) for x in header]

        #: if we are passing a single item, not an array of sets, don't map over it.
        #: this is when we have a station and not a set of results
        if not isinstance(rows[0], tuple):
            return dict(list(zip(header, rows)))

        return [dict(list(zip(header, x))) for x in rows]

    def _get_file_name_without_extension(self, file_path):
        '''Given a filename with an extension, the file name is returned without the extension.'''

        return splitext(basename(file_path))[0]

    def _add_sample_index(self, filepath):
        '''Add an index to ActivityIdentifier field for the table matching the file'''
        with sqlite3.connect(self.TEMPDB) as conn:
            conn.cursor().execute(self.sql['create_index'].format(basename(filepath)[:-4]))

    def _format_url(self, template, source, last_updated, today=None):
        date_format = '%m-%d-%Y'
        lo = dateparser(last_updated).strftime(date_format)
        hi = datetime.now().strftime(date_format)

        if today:
            hi = dateparser(today).strftime(date_format)

        return template.format(source, lo, hi)

    def _group_rows_by_id(self, cursor, config=None):
        '''groups samples by SampleId as they would be formatted by querycsv
        cursor: generator
        config: an optional config to setup the etl. mainly for testing

        returns a dicionary with sample_id's as the key, with a list of rows as values
        '''
        unique_sample_ids = {}

        if not cursor:
            return

        header = next(cursor)

        for row in cursor:
            row = self._etl_column_names(row, config or self.result_config, header=header)

            sample_id = row['SampleId']
            if sample_id in unique_sample_ids:
                unique_sample_ids[sample_id] += [row]
            else:
                unique_sample_ids[sample_id] = [row]

        return unique_sample_ids

    def _find_new_station_ids(self, rows):
        '''gets the station ids to process from the rows
        rows: {sampleId: list(results)} the sorted results from an update operation

        returns a set of station ids
        '''

        if not rows or len(rows) < 1:
            return []

        unique_station_ids = set([])

        for results in list(rows.values()):
            for result in results:
                station_id = result['StationId']

                unique_station_ids.add(station_id)

        unique_station_ids = self._get_unique_station_ids(list(unique_station_ids))

        return [id[0] for id in unique_station_ids]

    def _extract_stations_by_id(self, cursor, station_ids, header):
        '''loops over a cursor of stations and returns the stations that have an id
        in station_ids
        cursor: a generator of stations. they can be ETL'd or not
        station_ids: list(stationids)
        header: list(column names)

        returns a list of stations
        '''
        if not cursor:
            return

        stations_to_insert = []

        for row in cursor:
            #: have to etl row to check station id
            row = self._etl_column_names(row, self.station_config, header=header)

            if row['StationId'] in station_ids:
                stations_to_insert.append(row)

        return stations_to_insert

    def _remove_existing_wqx_station_ids(self, station_ids):
        '''takes a list of station ids and finds the wqx ids. It strips the wqx
        and then queries the database to see if the stripped id is already in the db.
        Removing any station id's that are already in the db.
        station_ids: list('station_id')

        returns a list of station_ids
        '''
        #: pull out wqx ids from Stations
        wqxs = set([id for id in [x for x in station_ids if self.wqx_re.search(x)]])

        #: if there are no wqx's then return
        if len(wqxs) < 1:
            return station_ids

        #: remove the wqx because the database does not store that value
        stripped_wqx = [re.sub(self.wqx_re, '-', id) for id in wqxs]

        #: get back the unique id's from the wqx subset
        uniques = set([stripped[0] for stripped in self._get_unique_station_ids(stripped_wqx)])

        #: remove the ids that are not in unique
        for id in station_ids:
            if re.sub(self.wqx_re, '-', id) not in uniques:
                station_ids.remove(id)

        return station_ids

    def _get_unique_sample_ids(self, sample_ids):
        if not hasattr(self, 'cursor') or not self.cursor:
            self.cursor = self.cursor_factory(self.db['connection_string'])

        if len(sample_ids) == 0:
            return None

        statement = self.sql['new_results'].format(','.join(sample_ids))

        self.cursor.execute(statement)

        return self.cursor.fetchall()

    def _remove_existing_results(self, results):
        sample_ids = ['(\'{}\')'.format(re.sub('\'', '\'\'', sample_id)) for sample_id in list(results.keys())]

        unique_sample_ids = self._get_unique_sample_ids(sample_ids)
        #: flatten list
        unique_sample_ids = [item for iter_ in unique_sample_ids for item in iter_]

        return {key: results[key] for key in results if key in unique_sample_ids}

class SdwisProgram(Program):
    '''class for handling sdwis database rows'''

    datasource = 'SDWIS'

    sql = {
        'new_results': ('SELECT * FROM (VALUES{}) AS t(SampleDate,SampleId,Param) WHERE NOT EXISTS('
                        'SELECT 1 FROM [UGSWaterChemistry].[ugswaterchemistry].[Results] WHERE [SampleId] = t.SampleId AND '
                        '[SampleDate] = t.SampleDate AND [Param] = t.Param)'),
        'specific-result': '''SELECT
            UTV80.TSASAMPL.COLLLECTION_END_DT AS "SampleDate",
            UTV80.TSASAMPL.LAB_ASGND_ID_NUM AS "SampleId",
            UTV80.TSAANLYT.NAME AS "Param",
            UTV80.TINWSF.TINWSF_IS_NUMBER AS "StationId",
            UTV80.TSASAR.ANALYSIS_START_DT AS "AnalysisDate",
            UTV80.TSALAB.LAB_ID_NUMBER AS "LabName",
            UTV80.TSASAR.DETECTN_LIMIT_NUM AS "MDL",
            UTV80.TSASAR.DETECTN_LIM_UOM_CD AS "MDLUnit",
            UTV80.TINWSYS.TINWSYS_IS_NUMBER AS "OrgId",
            UTV80.TINWSYS.NAME AS "OrgName",
            UTV80.TSASAR.CONCENTRATION_MSR AS "ResultValue",
            UTV80.TSASAMPL.COLLCTN_END_TIME AS "SampleTime",
            UTV80.TINWSF.TYPE_CODE AS "SampMedia",
            UTV80.TSASAR.UOM_CODE AS "Unit",
            UTV80.TINLOC.LATITUDE_MEASURE AS "Lat_Y",
            UTV80.TINLOC.LONGITUDE_MEASURE AS "Lon_X",
            UTV80.TSAANLYT.CAS_REGISTRY_NUM AS "CAS_Reg",
            UTV80.TSASAR.TSASAR_IS_NUMBER AS "IdNum"
            FROM UTV80.TINWSF
            JOIN UTV80.TINWSYS ON
            UTV80.TINWSF.TINWSYS_IS_NUMBER = UTV80.TINWSYS.TINWSYS_IS_NUMBER
            JOIN UTV80.TINLOC ON
            UTV80.TINWSF.TINWSF_IS_NUMBER = UTV80.TINLOC.TINWSF_IS_NUMBER
            JOIN UTV80.TSASMPPT ON
            UTV80.TINWSF.TINWSF_IS_NUMBER = UTV80.TSASMPPT.TINWSF0IS_NUMBER
            JOIN UTV80.TSASAMPL ON
            UTV80.TSASMPPT.TSASMPPT_IS_NUMBER = UTV80.TSASAMPL.TSASMPPT_IS_NUMBER
            JOIN UTV80.TSASAR ON
            UTV80.TSASAMPL.TSASAMPL_IS_NUMBER = UTV80.TSASAR.TSASAMPL_IS_NUMBER
            JOIN UTV80.TSAANLYT ON
            UTV80.TSASAR.TSAANLYT_IS_NUMBER = UTV80.TSAANLYT.TSAANLYT_IS_NUMBER
            JOIN UTV80.TSALAB ON
            UTV80.TSASAMPL.TSALAB_IS_NUMBER = UTV80.TSALAB.TSALAB_IS_NUMBER
            WHERE UTV80.TINWSF.TYPE_CODE in ('SP', 'WL', 'IN', 'SS') AND
                    UTV80.TSASAR.CONCENTRATION_MSR IS NOT NULL AND
                    UTV80.TSASAMPL.COLLLECTION_END_DT = TO_DATE('{}','YYYY-MM-DD') AND
                    UTV80.TSASAMPL.LAB_ASGND_ID_NUM = ? AND
                    UTV80.TSAANLYT.NAME = ?;''',
        'result': '''SELECT
            UTV80.TSASAMPL.COLLLECTION_END_DT AS "SampleDate",
            UTV80.TSASAMPL.LAB_ASGND_ID_NUM AS "SampleId",
            UTV80.TSAANLYT.NAME AS "Param",
            UTV80.TINWSF.TINWSF_IS_NUMBER AS "StationId",
            UTV80.TSASAR.ANALYSIS_START_DT AS "AnalysisDate",
            UTV80.TSALAB.LAB_ID_NUMBER AS "LabName",
            UTV80.TSASAR.DETECTN_LIMIT_NUM AS "MDL",
            UTV80.TSASAR.DETECTN_LIM_UOM_CD AS "MDLUnit",
            UTV80.TINWSYS.TINWSYS_IS_NUMBER AS "OrgId",
            UTV80.TINWSYS.NAME AS "OrgName",
            UTV80.TSASAR.CONCENTRATION_MSR AS "ResultValue",
            UTV80.TSASAMPL.COLLCTN_END_TIME AS "SampleTime",
            UTV80.TINWSF.TYPE_CODE AS "SampMedia",
            UTV80.TSASAR.UOM_CODE AS "Unit",
            UTV80.TINLOC.LATITUDE_MEASURE AS "Lat_Y",
            UTV80.TINLOC.LONGITUDE_MEASURE AS "Lon_X",
            UTV80.TSAANLYT.CAS_REGISTRY_NUM AS "CAS_Reg",
            UTV80.TSASAR.TSASAR_IS_NUMBER AS "IdNum"
            FROM UTV80.TINWSF
            JOIN UTV80.TINWSYS ON
            UTV80.TINWSF.TINWSYS_IS_NUMBER = UTV80.TINWSYS.TINWSYS_IS_NUMBER
            JOIN UTV80.TINLOC ON
            UTV80.TINWSF.TINWSF_IS_NUMBER = UTV80.TINLOC.TINWSF_IS_NUMBER
            JOIN UTV80.TSASMPPT ON
            UTV80.TINWSF.TINWSF_IS_NUMBER = UTV80.TSASMPPT.TINWSF0IS_NUMBER
            JOIN UTV80.TSASAMPL ON
            UTV80.TSASMPPT.TSASMPPT_IS_NUMBER = UTV80.TSASAMPL.TSASMPPT_IS_NUMBER
            JOIN UTV80.TSASAR ON
            UTV80.TSASAMPL.TSASAMPL_IS_NUMBER = UTV80.TSASAR.TSASAMPL_IS_NUMBER
            JOIN UTV80.TSAANLYT ON
            UTV80.TSASAR.TSAANLYT_IS_NUMBER = UTV80.TSAANLYT.TSAANLYT_IS_NUMBER
            JOIN UTV80.TSALAB ON
            UTV80.TSASAMPL.TSALAB_IS_NUMBER = UTV80.TSALAB.TSALAB_IS_NUMBER
            WHERE UTV80.TINWSF.TYPE_CODE in ('SP', 'WL', 'IN', 'SS') AND
                    UTV80.TSASAR.CONCENTRATION_MSR IS NOT NULL''',
        'date_clause': ' AND UTV80.TSASAMPL.COLLLECTION_END_DT > TO_DATE(\'{}\',\'YYYY-MM-DD\')',
        'unique_sample_ids': '''SELECT DISTINCT UTV80.TSASAMPL.LAB_ASGND_ID_NUM AS "SampleId"
            FROM UTV80.TINWSF
            JOIN UTV80.TSASMPPT ON
                UTV80.TINWSF.TINWSF_IS_NUMBER = UTV80.TSASMPPT.TINWSF0IS_NUMBER
            JOIN UTV80.TSASAMPL ON
                UTV80.TSASMPPT.TSASMPPT_IS_NUMBER = UTV80.TSASAMPL.TSASMPPT_IS_NUMBER
            JOIN UTV80.TSASAR ON
                UTV80.TSASAMPL.TSASAMPL_IS_NUMBER = UTV80.TSASAR.TSASAMPL_IS_NUMBER
            WHERE UTV80.TINWSF.TYPE_CODE in ('SP', 'WL', 'IN', 'SS') AND
                    UTV80.TSASAR.CONCENTRATION_MSR IS NOT NULL''',
        'sample_id': ' AND UTV80.TSASAMPL.LAB_ASGND_ID_NUM = ?',
        'station': '''SELECT
            UTV80.TINWSYS.TINWSYS_IS_NUMBER AS "OrgId",
            UTV80.TINWSYS.NAME AS "OrgName",
            UTV80.TINWSF.TINWSF_IS_NUMBER AS "StationId",
            UTV80.TINWSF.NAME AS "StationName",
            UTV80.TINWSF.TYPE_CODE AS "StationType",
            UTV80.TINLOC.LONGITUDE_MEASURE AS "Lon_X",
            UTV80.TINLOC.LATITUDE_MEASURE AS "Lat_Y",
            UTV80.TINLOC.HORIZ_ACCURACY_MSR AS "HorAcc",
            UTV80.TINLOC.HZ_COLLECT_METH_CD AS "HorCollMeth",
            UTV80.TINLOC.HORIZ_REF_DATUM_CD AS "HorRef",
            UTV80.TINLOC.VERTICAL_MEASURE AS "Elev",
            UTV80.TINLOC.VERT_ACCURACY_MSR AS "ElevAcc",
            UTV80.TINLOC.VER_COL_METH_CD AS "ElevMeth",
            UTV80.TINLOC.VERT_REF_DATUM_CD AS "ElevRef",
            MAX(UTV80.TINWLCAS.BOTTOM_DEPTH_MSR) AS "Depth",
            UTV80.TINWLCAS.BOTTOM_DP_MSR_UOM AS "DepthUnit"
            FROM UTV80.TINWSF
            JOIN UTV80.TINWSYS ON
                UTV80.TINWSF.TINWSYS_IS_NUMBER = UTV80.TINWSYS.TINWSYS_IS_NUMBER
            JOIN UTV80.TINLOC ON
                UTV80.TINWSF.TINWSF_IS_NUMBER = UTV80.TINLOC.TINWSF_IS_NUMBER
            LEFT JOIN UTV80.TINWLCAS ON
                UTV80.TINWSF.TINWSF_IS_NUMBER = UTV80.TINWLCAS.TINWSF_IS_NUMBER
            WHERE (UTV80.TINWSF.TYPE_CODE = 'SP' Or
                    UTV80.TINWSF.TYPE_CODE = 'WL' Or
                    UTV80.TINWSF.TYPE_CODE = 'IN' Or
                    UTV80.TINWSF.TYPE_CODE = 'SS') AND
                    UTV80.TINLOC.LATITUDE_MEASURE != 0 {}
            GROUP BY UTV80.TINWSF.TINWSF_IS_NUMBER,
                UTV80.TINWSF.NAME,
                UTV80.TINWSF.TYPE_CODE,
                UTV80.TINWSYS.TINWSYS_IS_NUMBER,
                UTV80.TINWSYS.NAME,
                UTV80.TINLOC.LATITUDE_MEASURE,
                UTV80.TINLOC.LONGITUDE_MEASURE,
                UTV80.TINLOC.SRC_MAP_SCALE_NUM,
                UTV80.TINLOC.HORIZ_ACCURACY_MSR,
                UTV80.TINLOC.HZ_COLLECT_METH_CD,
                UTV80.TINLOC.HORIZ_REF_DATUM_CD,
                UTV80.TINLOC.VERTICAL_MEASURE,
                UTV80.TINLOC.VERT_ACCURACY_MSR,
                UTV80.TINLOC.VER_COL_METH_CD,
                UTV80.TINLOC.VERT_REF_DATUM_CD,
                UTV80.TINWLCAS.BOTTOM_DEPTH_MSR,
                UTV80.TINWLCAS.BOTTOM_DP_MSR_UOM''',
        'station_id': 'AND UTV80.TINWSF.TINWSF_IS_NUMBER IN ({})'
    }

    def __init__(self,
                 logger,
                 db,
                 update,
                 source,
                 secrets=None,
                 sql_statements={},
                 update_row=None,
                 insert_rows=None,
                 cursor_factory=None,
                 arcpy=None):
        '''create a new SDWIS program
        db - the connection string for the database to seed
        source - the connection information to the sdwis database
        update - boolean value whether we are seeding or updating
        secrets - sdwis database connection information
        sql_statements - common sql statements for inserting into stations and results
        update_row - the function to reproject a point and set the DataSource
        insert_row - the function to batch insert rows
        cursor_factory - the function to create pyodbc connection_string
        arcpy - ignored. for gdb programs only
        '''
        self.logger = logger
        self.db = db
        self.source_db = secrets
        self._update_row = update_row
        self._insert_rows = insert_rows
        self.sql.update(sql_statements)
        self.source_cursor = cursor_factory(secrets['connection_string'])
        self.cursor_factory = cursor_factory

    def seed(self):
        try:
            self.logger.info('seeding {} stations...'.format(self.datasource))

            self._seed_stations(self.source_cursor.execute(self.sql['station'].format('')), schema.station)

            self.logger.info('seeding {} stations done.'.format(self.datasource))
            self.logger.info('seeding {} results...'.format(self.datasource))

            self._seed_results(self.source_cursor.execute(self.sql['unique_sample_ids']))

            self.logger.info('seeding {} results done.'.format(self.datasource))
        finally:
            if hasattr(self, 'source_cursor'):
                del self.source_cursor
            if hasattr(self, 'second_source_cursor'):
                del self.second_source_cursor
            if hasattr(self, 'cursor'):
                del self.cursor

    def update(self):
        self.logger.info('updating {} stations...'.format(self.datasource))

        try:
            #: query for new Results
            last_updated = self._get_most_recent_result_date(self.datasource)

            if not last_updated:
                raise Exception('No last updated date. You should seed some data first.')

            self.logger.info('fetching records after {}'.format(last_updated))
            last_updated = last_updated.split(' ')[0]
            query = self.sql['result'] + self.sql['date_clause'].format(last_updated)

            #: group new results by SampleDate + SampleId + Param
            new_results = self._group_rows_by_id(self.source_cursor.execute(query))

            if len(new_results) <= 0:
                self.logger.info('No new results to update. Quitting.')
                return

            new_results = self._remove_existing_results(new_results)

            #: find the station ids from the new results that aren't in the database
            new_station_ids = self._find_new_station_ids(new_results)

            if new_station_ids and len(new_station_ids) > 0:
                self.logger.info('adding {} new stations...'.format(len(new_station_ids)))
                end = 500
                while new_station_ids:
                    stations = new_station_ids[0:end]
                    query = self.sql['station'].format(self.sql['station_id']).format(','.join(stations))
                    self._seed_stations(self.source_cursor.execute(query), schema.station)

                    #: remove stations already inserted
                    del new_station_ids[0:end]

                self.logger.info('updating {} stations done.'.format(self.datasource))

            self.logger.info('adding {} new results...'.format(len(list(new_results.keys()))))

            self._seed_results(list(new_results.keys()))

            self.logger.info('updating {} results done.'.format(self.datasource))
        finally:
            if hasattr(self, 'source_cursor'):
                del self.source_cursor
            if hasattr(self, 'second_source_cursor'):
                del self.second_source_cursor
            if hasattr(self, 'cursor'):
                del self.cursor

    def _seed_stations(self, rows, config=None):
        #: rows cursor()
        #: config=The schema names for the rows
        stations = []

        for row in rows:
            #: zip column names with values
            row = self._zip_column_names(row)
            #: add shape column so row gets sql shape added
            row['Shape'] = None

            #: cast
            row = Caster.cast(row, schema.station)

            #: set datasource, reproject and update shape
            row = self._update_row(row, self.datasource)

            #: normalize data including stripping _WXP etc
            row = Normalizer.normalize_station(row)

            #: reorder and filter out any fields not in the schema
            row = Normalizer.reorder_filter(row, schema.station)

            #: have to generate sql manually because of quoting on spatial WKT
            row = Caster.cast_for_sql(row)

            #: store row for later
            stations.append(list(row.values()))

        if not hasattr(self, 'cursor') or not self.cursor:
            self.cursor = self.cursor_factory(self.db['connection_string'])

        #: insert stations
        self._insert_rows(stations, self.sql['station_insert'], self.cursor)

    def _seed_results(self, results_or_result_key):
        '''
        given a cursor or a list of keys, query for each one and insert it.
        '''
        for sample_id in results_or_result_key:
            #: sample id for seed or sample_date, sample_id, param key for update
            samples = self._get_samples_for_id(sample_id)

            #: cast
            samples = [Caster.cast(self._zip_column_names(sample), schema.result) for sample in samples]

            #: set datasource and spatial information
            samples = [self._update_row(sample, self.datasource) for sample in samples]

            #: normalize chemical names and units
            samples = list(map(Normalizer.normalize_sample, samples))

            #: create charge balance rows from sample
            charge_balances = ChargeBalancer.get_charge_balance(samples)

            samples.extend(charge_balances)

            #: reorder and filter out any fields not in the schema
            samples = list(map(partial(Normalizer.reorder_filter, schema=schema.result), samples))

            samples = list(map(Caster.cast_for_sql, samples))

            rows = [list(sample.values()) for sample in samples]

            if not hasattr(self, 'cursor') or not self.cursor:
                self.cursor = self.cursor_factory(self.db['connection_string'])

            #: TODO determine if this should this be batched in sets bigger than just a sample set?
            self._insert_rows(rows, self.sql['result_insert'], self.cursor)

    def _get_samples_for_id(self, sample_id_or_key):
        #: updating
        try:
            key = sample_id_or_key.split('{-}')

            sample_date = key[0]
            sample_id = key[1]
            param = key[2]

            return self.source_cursor.execute(self.sql['specific-result'].format(sample_date), sample_id, param)
        except AttributeError:
            #: 'pyodbc.Row' object has no attribute 'split'
            pass

        #: seeding
        #: introduce another cursor so that seeding can continue to iterate over it's buffer
        if not hasattr(self, 'second_source_cursor') or not self.second_source_cursor:
            self.second_source_cursor = self.cursor_factory(self.source_db['connection_string'])

        return self.second_source_cursor.execute(self.sql['result'] + self.sql['sample_id'], sample_id_or_key[0])

    def _zip_column_names(self, row):
        '''Given a set, return a dictionary with field names'''

        #: we have an already etl'd station. skip it.
        if type(row) is dict:
            return row

        if len(row) == 0:
            return None

        header = [cd[0] for cd in row.cursor_description]

        return dict(list(zip(header, list(row))))

    def _find_new_station_ids(self, rows):
        '''gets the station ids to process from the rows
        rows: {key: list(stationId)}

        returns a set of station ids
        '''

        if not rows or len(rows) < 1:
            return []

        unique_station_ids = set([])

        for results in list(rows.values()):
            for result in results:
                station_id = result

                unique_station_ids.add(station_id)

        unique_station_ids = self._get_unique_station_ids(list(unique_station_ids))

        return [id[0] for id in unique_station_ids]

    def _group_rows_by_id(self, cursor):
        '''groups results by SampleDate + SampleId + Param
        cursor: generator over a sql query for results

        returns a dicionary with SampleDate + SampleId + Param as the key, with a list of station_ids
        '''
        unique_sample_ids = {}

        if not cursor:
            return

        for row in cursor:
            sample_date = row[0]
            sample_id = row[1]
            param = row[2]
            station_id = row[3]
            key = '{}{{-}}{}{{-}}{}'.format(sample_date.strftime('%Y-%m-%d'), sample_id, param)
            if key in unique_sample_ids:
                unique_sample_ids[key] += [station_id]
            else:
                unique_sample_ids[key] = [station_id]

        return unique_sample_ids

    def _remove_existing_results(self, results):
        sample_ids = []
        for key in list(results.keys()):
            key = key.split('{-}')

            sample_date = key[0]
            sample_id = key[1]
            param = key[2]
            sample_ids.append('(\'{}\',\'{}\',\'{}\')'.format(sample_date,
                                                              re.sub('\'', '\'\'', sample_id),
                                                              re.sub('\'', '\'\'', param)))

        unique_sample_ids = self._get_unique_sample_ids(sample_ids)

        #: munge back to stationid, date, param key
        unique_sample_ids = ['{}{{-}}{}{{-}}{}'.format(item[0], item[1], item[2]) for item in unique_sample_ids]

        return {key: results[key] for key in results if key in unique_sample_ids}

    def _get_unique_sample_ids(self, sample_ids):
        if not hasattr(self, 'cursor') or not self.cursor:
            self.cursor = self.cursor_factory(self.db['connection_string'])

        if len(sample_ids) == 0:
            return None

        statement = self.sql['new_results'].format(','.join(sample_ids))

        return self.cursor.execute(statement)

class GdbProgram(object):
    #: sql queries
    sql = {}

    def __init__(self,
                 logger,
                 db,
                 update,
                 source=None,
                 secrets=None,
                 sql_statements={},
                 update_row=None,
                 insert_rows=None,
                 cursor_factory=None,
                 arcpy=None):
        '''create a new DOGM program
        db - the connection string for the database to seed
        update - boolean value whether we are seeding or updating
        source - the path on disk to find the parent gdb folder to ETL
        secrets - ignored. for sdwis only
        sql_statements - common sql statements for inserting into stations and results
        update_row - the function to reproject a point and set the DataSource
        insert_row - the function to batch insert rows
        cursor_factory - the function to create pyodbc connection_string
        arcpy - the arcpy module
        '''
        self.logger = logger
        self.db = db
        self._update_row = update_row
        self._insert_rows = insert_rows
        self.arcpy = arcpy
        self.cursor_factory = cursor_factory
        self.sql.update(sql_statements)

        #: check that source exists wqp/results and wqp/stations
        parent_folder = join(source, self.gdb)

        if not isdir(parent_folder):
            raise Exception('Pass in a location to the parent folder that contains DOGM. {}'.format(parent_folder))

        arcpy.env.workspace = parent_folder
        if not arcpy.Exists(self.result_table) or not arcpy.Exists(self.station_table):
            raise Exception('The DOGM geodatabase is missing the feature class {} or table.'.format(self.station_table,
                                                                                self.result_table))

        try:
            arcpy.RemoveIndex_management(self.result_table, index_name='result_query')
        except Exception:
            self.logger.error('There was an issue removing an index.')

        try:
            arcpy.AddIndex_management(in_table=self.result_table,
                                      fields='SampleId',
                                      index_name='result_query')
        except Exception:
            self.logger.error('There was an issue adding an index. Inserting may be slower or removing the existing index failed.')

    def seed(self):
        self._seed('seeding')

    def update(self):
        self._seed('updating')

    def _seed(self, what):
        #: location - the path to the table data
        #: fields - the fields form the data to pull
        self.logger.info('{} {} stations...'.format(what, self.datasource))

        #: get subset of dogm fields
        station_fields = self._get_field_instersection(schema.station, self.arcpy.ListFields(self.station_table))

        #: use geodatabase shape
        station_fields.remove('Shape')
        station_fields.append('Shape@XY')

        try:
            self.source_cursor = self.arcpy.da.SearchCursor(self.station_table, field_names=station_fields,
                                                            where_clause='1=1')
            self._seed_stations(self.source_cursor, station_fields)

            self.logger.info('{} {} stations done.'.format(what, self.datasource))
            self.logger.info('{} {} results...'.format(what, self.datasource))

            self.source_cursor = self.arcpy.da.SearchCursor(self.result_table,
                                                            field_names=['SampleId'],
                                                            sql_clause=('DISTINCT', None))

            self._seed_results(self.source_cursor)

            self.logger.info('{} {} results done.'.format(what, self.datasource))
        finally:
            if hasattr(self, 'source_cursor'):
                del self.source_cursor
            if hasattr(self, 'second_source_cursor'):
                del self.second_source_cursor
            if hasattr(self, 'cursor'):
                del self.cursor

    def _seed_stations(self, rows, config=None):
        #: rows cursor()
        #: config=The schema names for the rows
        stations = []

        for row in rows:
            #: zip column names with values
            row = self._zip_column_names(row, config)

            #: cast
            row = Caster.cast(row, schema.station)

            #: set datasource, reproject and update shape
            row = self._update_row(row, self.datasource)

            if row['Shape'] is None:
                self.logger.warn('Skipping row because it has an invalid shape.')
                self.logger.debug(row)

                continue

            #: normalize data including stripping _WXP etc
            row = Normalizer.normalize_station(row)

            #: reorder and filter out any fields not in the schema
            row = Normalizer.reorder_filter(row, schema.station)

            #: have to generate sql manually because of quoting on spatial WKT
            row = Caster.cast_for_sql(row)

            #: store row for later
            stations.append(list(row.values()))

        if not hasattr(self, 'cursor') or not self.cursor:
            self.cursor = self.cursor_factory(self.db['connection_string'])

        #: insert stations
        self._insert_rows(stations, self.sql['station_insert'], self.cursor)

    def _seed_results(self, sample_ids):
        result_fields = self._get_field_instersection(schema.result, self.arcpy.ListFields(self.result_table))

        for sample_id in sample_ids:
            sample_id = sample_id[0]

            try:
                samples = self._get_samples_for_id(sample_id, result_fields)
            finally:
                if hasattr(self, 'second_source_cursor'):
                    del self.second_source_cursor

            #: cast
            samples = [Caster.cast(self._zip_column_names(sample, result_fields), schema.result) for sample in samples]

            #: set datasource and spatial information
            samples = [self._update_row(sample, self.datasource) for sample in samples]

            #: normalize chemical names and units
            samples = list(map(Normalizer.normalize_sample, samples))

            #: create charge balance rows from sample
            charge_balances = ChargeBalancer.get_charge_balance(samples)

            samples.extend(charge_balances)

            #: reorder and filter out any fields not in the schema
            samples = list(map(partial(Normalizer.reorder_filter, schema=schema.result), samples))

            samples = list(map(Caster.cast_for_sql, samples))

            rows = [list(sample.values()) for sample in samples]

            if not hasattr(self, 'cursor') or not self.cursor:
                self.cursor = self.cursor_factory(self.db['connection_string'])

            #: TODO determine if this should this be batched in sets bigger than just a sample set?
            self._insert_rows(rows, self.sql['result_insert'], self.cursor)

    def _get_samples_for_id(self, sample_id, result_fields):
        #: introduce another cursor so that seeding can continue to iterate over it's buffer
        self.second_source_cursor = self.arcpy.da.SearchCursor(self.result_table,
                                                               field_names=result_fields,
                                                               where_clause='SampleId=\'{}\''.format(sample_id))

        return self.second_source_cursor

    def _get_field_instersection(self, ugs_schema, table_fields):
        station_fields = set([str(x.name) for x in table_fields])
        schema_fields = set(ugs_schema.keys())

        return list(station_fields & schema_fields)

    def _zip_column_names(self, row, fields):
        '''Given a set, return a dictionary with field names'''

        #: we have an already etl'd station. skip it.
        if type(row) is dict:
            return row

        if len(row) == 0:
            return None

        return dict(list(zip(fields, list(row))))

class DogmProgram(GdbProgram):
    datasource = 'UDOGM'
    #: location to dogm gdb
    gdb = join('DOGM', 'DOGM_AGRC.gdb')

    #: results table name
    result_table = 'DOGM_RESULT'

    #: stations feature class name
    station_table = 'DOGM_STATION'

class UdwrProgram(GdbProgram):
    datasource = 'UDWR'
    #: location to dogm gdb
    gdb = join('UDWR', 'UDWR_AGRC.gdb')

    #: results table name
    result_table = 'UDWR_RESULTS'

    #: stations feature class name
    station_table = 'UDWR_STATION'

class UgsProgram(GdbProgram):
    datasource = 'UGS'
    #: location to dogm gdb
    gdb = join('UGS', 'UGS_AGRC.gdb')

    #: results table name
    result_table = 'RESULTS'

    #: stations feature class name
    station_table = 'STATIONS'


from dataclasses import dataclass
from os import getcwd
from pathlib import Path

import click
from pytorch_lightning import Trainer
from pytorch_lightning.loggers import TensorBoardLogger

from ad_denoise.lightning_modules import LightningModuleConfig
from ad_denoise.utils import load_config

from . import __version__

@click.group(invoke_without_command=True)
@click.version_option(version=__version__, message="%(version)s")
@click.pass_context
def main(ctx: click.Context) -> None:
    """The main command line interface entry point.

    Args:
        ctx: The click context.
    """
    if ctx.invoked_subcommand is None:
        click.echo(main.get_help(ctx))

@dataclass
class TrainConfig:
    """A configuration schema for network training."""

    name: str
    model: LightningModuleConfig
    max_epochs: int

@main.command(help="Train a model on the given datasets")
@click.argument("config_file", type=click.Path(exists=True, dir_okay=False))
def train(config_file: Path) -> None:  # noqa: D103
    config = load_config(config_file, TrainConfig)
    logger = TensorBoardLogger(str(Path(getcwd()).joinpath("logs")), config.name)
    trainer = Trainer(
        max_epochs=config.max_epochs,
        log_every_n_steps=1,
        accelerator="auto",
        logger=logger,
    )
    trainer.fit(config.model())


import datetime
import influxdb
import multiprocessing
import time
import orchestration.timing
import orchestration.queues
import os

#=======================================================================================================================
class Counter(object):
    __slots__ = ['__name', '__group', '__unit', '__time' ,'value']

#-----------------------------------------------------------------------------------------------------------------------
    def __init__(self, name, group, unit):
        """ C'tor
        @param name: Counter name
        @param group: The contextual group in which this counter belongs to (e.g. disk-red-bytes and disk-written-bytes
                      belongs to group disk-io)
        @param unit: The unit name in which this counter is counting for (e.g. Events handler name)
        """

        self.value = 0
        self.__time = -1
        self.__name = name
        self.__group = group
        self.__unit = unit

#-----------------------------------------------------------------------------------------------------------------------
    @property
    def name(self):
        """ The counter's name
        """

        return self.__name

#-----------------------------------------------------------------------------------------------------------------------
    @property
    def group(self):
        """ The counter's group name
        """

        return self.__group

#-----------------------------------------------------------------------------------------------------------------------
    @property
    def unit(self):
        """ The counter's unit name
        """

        return self.__unit

#-----------------------------------------------------------------------------------------------------------------------
    def at(self, time):
        """ Set a time for the counter so next value set will be regarded as set in this time
        @param time: Unix time
        @return: This object with the time set for chaining value set
        """

        self.__time = time

        return self

#-----------------------------------------------------------------------------------------------------------------------
    @property
    def _time(self):
        """ Unix time in case at() value set was used, -1 if was never used
        """

        return self.__time

#-----------------------------------------------------------------------------------------------------------------------
    def __repr__(self): return "Counter(name=%s, group=%s, unit=%s, value=%s)" % (self.__name, self.__group,
                                                                                self.__unit, self.value)

#=======================================================================================================================
class CountersContainer(object):

#-----------------------------------------------------------------------------------------------------------------------
    def __init__(self, logger, id_, name, configuration):
        """ C'tor
        @param logger: Parent logger
        @param id_: Integer to identify the container with
        @param name: container name
        @param orchestration configuration
        @see configuration.Configuration
        """

        self.__id = id_
        self.__name = name
        self.__logger = logger.createAdapter(self)
        self.__configuration = configuration
        self.__counters = []
        self.__queue = orchestration.queues._QueueZmq(self.__logger, 0, "egress-stats-collector",
                                                      os.path.join(self.__configuration.runDirectoryPath,
                                                                   _Collector.UNIX_DOMAIN_SOCKET),
                                                      orchestration.queues._SerializerPickle())

#-----------------------------------------------------------------------------------------------------------------------
    def registerCounter(self, counter):
        """ Register a counter to this container
        @param counter: Counter to register
        @return: Registered counters
        """

        self.__logger.debug2("Registering counter: counter=%s", counter)
        self.__counters.append(counter)

        return counter

#-----------------------------------------------------------------------------------------------------------------------
    @property
    def _name(self):
        """ The counters container name
        """

        return self.__name

#-----------------------------------------------------------------------------------------------------------------------
    def _connectToCollector(self):
        """ Connect to counters collector
        """

        self.__logger.debug2("Connecting to collector")
        self.__queue.connect()

#-----------------------------------------------------------------------------------------------------------------------
    def _disconnectFromCollector(self):
        """ Disconnect from counters collector
        """

        self.__logger.debug2("Disconnecting from collector")
        self.__queue.disconnect()

#-----------------------------------------------------------------------------------------------------------------------
    def _iterCounters(self):
        """ Create an iterator for counters
        @return: Iterator for counters
        """

        for counter in self.__counters:
            yield counter

#-----------------------------------------------------------------------------------------------------------------------
    def _publish(self):
        """ Publish to collector
        """

        currentTime = int(time.time())

        if self.__counters:
            values = []
            times = []

            for counter in self.__counters:
                values.append(counter.value)
                times.append(currentTime if counter._time == -1 else counter._time)

            dataPoints = _DataPoints(self.__id, values, times)

            if self.__logger.isEnabledFor(orchestration.logger.DEBUG5):
                self.__logger.debug5("Sending data points to collector: dataPoints=%s", dataPoints)

            self.__queue.push(dataPoints)

#-----------------------------------------------------------------------------------------------------------------------
    def __repr__(self): return "CountersContainer(name=%s, id=%s)" % (self.__name, self.__id)

#=======================================================================================================================
class _DataPoints(object):
    __slots__ = ['countersContainerId', 'values', 'times']

#-----------------------------------------------------------------------------------------------------------------------
    def __init__(self, countersContainerId, values, times):
        """ C'tor
        @param countersContainerId: Id of counters container
        @param values: List of counters values
        @param times: List of counters times in Unix time
        """

        self.countersContainerId = countersContainerId
        self.values = values
        self.times = times

#-----------------------------------------------------------------------------------------------------------------------
    def __repr__(self): return "_DataPoints(countersContainerId=%s, values=%s)" % (self.countersContainerId,
                                                                                self.values)

#=======================================================================================================================
class _Collector(object):

    UNIX_DOMAIN_SOCKET = "stats.sock"
    INFLUXDB_TAG_INCEPTION = "inception"

#-----------------------------------------------------------------------------------------------------------------------
    def __init__(self, logger, loggerListener, configuration):
        """ C'tor
        @param logger: Parent logger
        @param loggerListener: Logger listener to connect to
        @param orchestration configuration
        @see configuration.Configuration
        """

        self.__logger = logger.createAdapter("StatsCollector")
        self.__loggerListener = loggerListener
        self.__configuration = configuration
        self.__shouldRun = multiprocessing.Value('b', False)
        self.__wasLaunched = multiprocessing.Value('b', False)
        self.__process = multiprocessing.Process(target = self.__processTarget, name="stats-collector")
        self.__startSemaphore = multiprocessing.Semaphore()
        self.__queue = orchestration.queues._QueueZmq(self.__logger, 0, "ingress-stats-collector",
                                                      os.path.join(self.__configuration.runDirectoryPath,
                                                                   self.UNIX_DOMAIN_SOCKET),
                                                      orchestration.queues._SerializerPickle())
        self.__countersContainers = []
        self.__influxdbDataPoints = []
        self.__influxdbClient = None
        self.__inception = datetime.datetime.now().strftime("%x-%X")

        if not self.__configuration.disableStats:
            splitData = self.__configuration.influxdb.split('@')
            self.__influxdbDatabaseName = splitData[0]
            self.__influxdbHost = splitData[1].split(':')[0]
            self.__influxdbPort = splitData[1].split(':')[1]
            self.__flush_Poller = orchestration.timing._Poller(self.__flush,
                                                            self.__configuration.statsCollectorInfluxdbPublishInterval)

#-----------------------------------------------------------------------------------------------------------------------
    @property
    def isAlive(self):
        """ Is the logger listener alive
        """

        return self.__configuration.disableStats or self.__process.is_alive()

#-----------------------------------------------------------------------------------------------------------------------
    def createCountersContainer(self, name):
        """ Create a counters container
        @param name: container name
        @return: New counters container
        """

        containerId = len(self.__countersContainers)
        container = CountersContainer(self.__logger,containerId, name, self.__configuration)
        self.__logger.debug2("Counters container created: countersContainer=%s", (container,))
        self.__countersContainers.append(container)

        return container

#-----------------------------------------------------------------------------------------------------------------------
    def start(self):
        """ Start the stats collector
        """
        if not self.__configuration.disableStats:
            self.__logger.debug2("Starting stats collector")
            self.__shouldRun.value = True
            self.__process.start()

            # Acquire first to negate the semaphore and then try again till it is released
            self.__startSemaphore.acquire()
            self.__startSemaphore.acquire(True)

            if not self.__wasLaunched.value:
                self.__logger.error("Failed to launch stats collector")

                raise orchestration.QuietError()

#-----------------------------------------------------------------------------------------------------------------------
    def stop(self):
        """ Stop the stats collector
        """

        if not self.__configuration.disableStats:

            if self.__wasLaunched.value:
                self.__logger.debug2("Stopping stats collector")
                self.__shouldRun.value = False
                self.__process.join()

#-----------------------------------------------------------------------------------------------------------------------
    def __connectToInfluxdb(self):
        """ Connect to influxdb
        """

        self.__logger.debug2("Connecting to influxdb: host=%s, port=%s, database=%s", self.__influxdbHost,
                             self.__influxdbPort, self.__influxdbDatabaseName)

        try:
            self.__influxdbClient = influxdb.InfluxDBClient(host = self.__influxdbHost, port = self.__influxdbPort,
                                                            database = self.__influxdbDatabaseName)

        except:
            self.__logger.error("Failed connecting to influxdb: host=%s, port=%s, database=%s",
                                self.__influxdbHost, self.__influxdbPort, self.__influxdbDatabaseName, e)

            raise orchestration.QuietError()

#-----------------------------------------------------------------------------------------------------------------------
    def __validateInfluxdbDatabase(self):
        """ Validate the existence of the database in influxdb
        """

        self.__logger.debug2("Making sure database exists in influxdb: host=%s, port=%s, database=%s",
                             self.__influxdbHost, self.__influxdbPort, self.__influxdbDatabaseName)

        try:
            results = self.__sendInfluxdbQuery("SHOW DATABASES")
            databaseFound = self.__influxdbDatabaseName in [str(x['name']) for x in results['databases']]

            if not databaseFound:
                self.__logger.debug2("Influxdb database was not found, creating: host=%s, port=%s, database=%s",
                                     self.__influxdbHost, self.__influxdbPort, self.__influxdbDatabaseName)

                self.__sendInfluxdbQuery("CREATE DATABASE \"%s\"" % self.__influxdbDatabaseName)

        except:
            self.__logger.error("Failed validating influxdb database: host=%s, port=%s, database=%s",
                                self.__influxdbHost, self.__influxdbPort, self.__influxdbDatabaseName)

            raise orchestration.QuietError()

#-----------------------------------------------------------------------------------------------------------------------
    def __retainInfluxdbDatabaseSize(self):
        """ Retain the history size of our measures
        """

        self.__logger.debug2("Retaining influxdb database size: host=%s, port=%s, database=%s",
                             self.__influxdbHost, self.__influxdbPort, self.__influxdbDatabaseName)

        try:
            results = self.__sendInfluxdbQuery("SHOW TAG VALUES ON \"%s\" WITH KEY = \"%s\"" %
                                             (self.__influxdbDatabaseName, self.INFLUXDB_TAG_INCEPTION))

            for measurementName in [str(x[0]) for x in results.keys()]:
                history = sorted([x['value'] for x in results[measurementName]])

                if len(history) > self.__configuration.influxdbMeasurementHistoryMaxSize:

                    size = len(history) - self.__configuration.influxdbMeasurementHistoryMaxSize

                    for serieInception in history[:size]:
                        self.__sendInfluxdbQuery("DROP SERIES FROM \"%s\" WHERE \"%s\" = '%s'" %
                                                 (measurementName, self.INFLUXDB_TAG_INCEPTION, serieInception))

        except:
            self.__logger.error("Failed retaining influxdb database size: host=%s, port=%s, database=%s",
                                self.__influxdbHost, self.__influxdbPort, self.__influxdbDatabaseName)

            raise orchestration.QuietError()

#-----------------------------------------------------------------------------------------------------------------------
    def __sendInfluxdbQuery(self, query):
        """ Send an influxdb query
        @param query: Query to send
        @return: influxdb.resultset object
        """

        self.__logger.debug3("Sending query to influxdb: host=%s, port=%s, database=%s, query=%s",
                             self.__influxdbHost, self.__influxdbPort, self.__influxdbDatabaseName, query)

        return self.__influxdbClient.query(query)

#-----------------------------------------------------------------------------------------------------------------------
    def __flush(self):
        """ Flush data points to influxdb
        """

        if self.__influxdbDataPoints:

            if self.__logger.isEnabledFor(orchestration.logger.DEBUG5):
                self.__logger.debug5("Flushing data points to influxdb: influxdbDataPoint=%s",
                                     self.__influxdbDataPoints)

            rc = self.__influxdbClient.write_points(self.__influxdbDataPoints, time_precision='s')

            if not rc:
                self.__logger.warning("Failed writing points to influxdb: points=%s",
                                      self.__influxdbDataPoints)

            self.__influxdbDataPoints = []

#-----------------------------------------------------------------------------------------------------------------------
    def __processTarget(self):
        """ Stats collector process target
        """

        try:
            # Connect to logger listener, bind to queue, connect to influxdb and then release start condition
            self.__logger.logger._connectToListener(self.__loggerListener.createHandler())
            self.__queue.bind()
            self.__connectToInfluxdb()
            self.__validateInfluxdbDatabase()
            self.__retainInfluxdbDatabaseSize()
            self.__wasLaunched.value = True
            self.__startSemaphore.release()
            self.__logger.debug2("Stats collector started")

            while self.__shouldRun.value:
                dataPoints = self.__queue.pop()

                if dataPoints is not None:

                    if self.__logger.isEnabledFor(orchestration.logger.DEBUG5):
                        self.__logger.debug5("Received data points: dataPoints=%s", dataPoints)

                    # Aggregating to batch
                    countersContainer = self.__countersContainers[dataPoints.countersContainerId]
                    index = 0

                    for counter in countersContainer._iterCounters():

                        influxdbDataPoint = {'measurement' : counter.group,
                                             'tags'        : {'source' : countersContainer._name,
                                                              'unit' : counter.unit,
                                                              self.INFLUXDB_TAG_INCEPTION : self.__inception},
                                             'time'        : dataPoints.times[index],
                                             'fields'      : {counter.name : dataPoints.values[index]}}

                        if self.__logger.isEnabledFor(orchestration.logger.DEBUG5):
                            self.__logger.debug5("Created influxdb data point: influxdbDataPoint=%s",
                                                 influxdbDataPoint)

                        self.__influxdbDataPoints.append(influxdbDataPoint)
                        index += 1

                else:
                    time.sleep(self.__configuration.statsCollectorIdleTime)

                self.__flush_Poller.poll()

        except KeyboardInterrupt:
            pass

        except orchestration.QuietError:
            pass

        except:
            self.__logger.exception("Unexpected exception")

        finally:

            if not self.__wasLaunched.value:
                self.__startSemaphore.release()

            self.__queue.unbind()
            self.__logger.debug2("Stats collector stopped")
            self.__logger.logger._disconnectFromListener()


import argparse
import os.path as op
import os
from braincoder.models import GaussianPRF2DWithHRF
from braincoder.hrf import CustomHRFModel, SPMHRFModel
from braincoder.optimize import ParameterFitter
from braincoder.utils import get_rsq
from gaze_prf.utils.data import get_data, get_prf_model, get_dm, get_grid_coordinates, get_masker, get_all_roi_labels
from braincoder.optimize import ParameterFitter
import pandas as pd
import numpy as np

from nilearn import image
from itertools import product

def main(subject, bids_folder):

    target_dir = op.join(bids_folder, 'derivatives', 'prf_predictive_accuracy', f'sub-{subject}',
                         'func', )

    if not op.exists(target_dir):
        os.makedirs(target_dir)

    dm = get_dm(bids_folder=bids_folder, task='GazeCenter',
                resize_factor=3.0)
    paradigm = dm.reshape((len(dm), -1))

    masker = get_masker(subject, roi='brainmask',
                        bids_folder=bids_folder)

    for task in ['AttendStim', 'AttendFix']:

        r2s_retinotopic = []
        r2s_spatiotopic = []

        for gaze in ['Left', 'Right', 'Center']:

            retinotopic_model = get_prf_model(
                subject, task=task, bids_folder=bids_folder)
            spatiotopic_model = get_prf_model(
                subject, task=task, bids_folder=bids_folder)

            if gaze == 'Left':
                spatiotopic_model.parameters['x'] += 4
            elif gaze == 'Right':
                spatiotopic_model.parameters['x'] -= 4

            data = get_data(subject, bids_folder=bids_folder,
                            # fullscreen=False,
                            task=task,
                            gaze=gaze,
                            masker=masker)

            print(data)

            data = data.groupby(['time']).mean()

            # Correct for baseline/amplitude effects
            for model in [retinotopic_model, spatiotopic_model]:
                optimizer = ParameterFitter(model, data, paradigm)
                model.parameters = optimizer.refine_baseline_and_amplitude(
                    model.parameters)

            retinotopic_predictions = retinotopic_model.predict(
                paradigm=paradigm)
            spatiotopic_predictions = spatiotopic_model.predict(
                paradigm=paradigm)

            print(retinotopic_predictions)

            r2_retinotopic = np.clip(
                get_rsq(data, retinotopic_predictions), 0, 1)
            r2_spatiotopic = np.clip(
                get_rsq(data, spatiotopic_predictions), 0, 1)

            r2s_retinotopic.append(r2_retinotopic)
            r2s_spatiotopic.append(r2_spatiotopic)

            gardner = (r2_retinotopic - r2_spatiotopic) / \
                (r2_retinotopic + r2_spatiotopic)

            masker.inverse_transform(retinotopic_predictions).to_filename(op.join(target_dir,
                                                                         f'sub-{subject}_task-{task}Gaze{gaze}_desc-retinotopic_predictions.nii.gz'))
            masker.inverse_transform(spatiotopic_predictions).to_filename(op.join(target_dir,
                                                                         f'sub-{subject}_task-{task}Gaze{gaze}_desc-spatiotopic_predictions.nii.gz'))

            masker.inverse_transform(data).to_filename(op.join(target_dir,
                                                                         f'sub-{subject}_task-{task}Gaze{gaze}_desc-average_bold.nii.gz'))

            masker.inverse_transform(r2_retinotopic).to_filename(op.join(target_dir,
                                                                         f'sub-{subject}_task-{task}Gaze{gaze}_desc-retinotopic_r2.nii.gz'))
            masker.inverse_transform(r2_spatiotopic).to_filename(op.join(target_dir,
                                                                         f'sub-{subject}_task-{task}Gaze{gaze}_desc-spatiotopic_r2.nii.gz'))
            masker.inverse_transform(gardner).to_filename(op.join(target_dir,
                                                                  f'sub-{subject}_task-{task}Gaze{gaze}_gardner.nii.gz'))

        r2s_retinotopic = np.mean(r2s_retinotopic[:2], 0)
        r2s_spatiotopic = np.mean(r2s_spatiotopic[:2], 0)
        gardner = (r2s_retinotopic - r2s_spatiotopic) / \
            (r2s_retinotopic + r2s_spatiotopic)

        df = []
        for d, label in zip([r2s_retinotopic, r2s_spatiotopic, gardner], ['desc-retinotopic_r2', 'desc-spatiotopic_r2',
                                                                          'gardner']):
            d = masker.inverse_transform(d)

            d.to_filename(
                op.join(target_dir, f'sub-{subject}_task-{task}_{label}.nii.gz'))
            df.append(d)

        df = image.concat_imgs(df)

        for roi in get_all_roi_labels():
            roi_masker = get_masker(subject, roi, bids_folder, 'both')

            p = roi_masker.fit_transform(df)
            p = pd.DataFrame(
                p.T, columns=['r2_retinotopic', 'r2_spatiotopic', 'gardner'])
            p.index.name = 'voxel'
            p.to_csv(op.join(
                target_dir, f'sub-{subject}_task-{task}_roi-{roi}_r2s.tsv'))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('subject', default=None)
    parser.add_argument(
        '--bids_folder', default='/tank/shared/2021/visual/pRFgazeMod/')

    args = parser.parse_args()

    main(args.subject, bids_folder=args.bids_folder)


# -*- coding: utf-8 -*-
"""
Created on Mon Apr  6 16:24:53 2020

@author: Mukul Kirti Verma
"""

# your task is to complete this function
# function should return index to the any valid peak element
def findMid(head):
    # Code here
    temp=head
    temp2=head
    if(head==None):

        return temp
    else:
        while(temp.next!=None):
            temp=temp.next
            if(temp.next!=None):
                temp=temp.next
            temp2=temp2.next
    return temp2

#{
#  Driver Code Starts
# Node Class
class node:
    def __init__(self, val):
        self.data = val
        self.next = None

# Linked List Class
class Linked_List:
    def __init__(self):
        self.head = None

    def insert(self, val):
        if self.head == None:
            self.head = node(val)
        else:
            new_node = node(val)
            temp = self.head
            while(temp.next):
                temp=temp.next
            temp.next = new_node

def createList(arr, n):
    lis = Linked_List()
    for i in range(n):
        lis.insert(arr[i])
    return lis.head

if __name__=='__main__':
    t = int(input())
    for i in range(t):
        n = int(input())
        arr = list(map(int, input().strip().split()))
        head = createList(arr, n)
        print(findMid(head).data)

# } Driver Code Ends


from __future__ import annotations

from typing import List, Optional, Tuple

import numpy as np
from sklearn.cluster import DBSCAN

from clientAis.ais.discreteV1.physics.PhysicsEstimator import PhysicsEstimator
from clientAis.ais.discreteV1.Plans.DiscreteMove import DiscreteMove
from clientAis.games.GameState import GameState, Singuity, Spawner
from utils import arrays

class DiscretePlayer:
    def __init__(
        self, isCurrentPlayer: bool, id: str, singuityCount: int, singuitiesMeanPosition: np.ndarray, singuitiesStd: float, singuitiesAverageHealth: float,
        inCluster: Optional[List[str]] = None
        ):
        self.isCurrentPlayer = isCurrentPlayer
        self.id = id
        self.singuityCount = singuityCount
        self.singuitiesMeanPosition = singuitiesMeanPosition
        self.singuitiesStd = singuitiesStd
        self.singuitiesAverageHealth = singuitiesAverageHealth
        self.inCluster = inCluster

    @staticmethod
    def fromGameState(gameState: GameState, playerId: str, currentPlayerId: str) -> DiscretePlayer:
        playerSinguities = [s for s in gameState.singuities if s.playerId == playerId]

        if len(playerSinguities) == 0:
            return DiscretePlayer(currentPlayerId == playerId, playerId, 0, np.zeros(2), PhysicsEstimator.getMinimumStd(0), 0, [])

        singuityPositions = [s.position for s in playerSinguities]

        singuityClusterIds = DBSCAN(eps=100, min_samples=25).fit_predict(singuityPositions)
        singuityClusters: List[List[Singuity]] = [[] for _ in range(np.max(singuityClusterIds) + 1)]

        for index, singuity in [(index, singuity) for index, singuity in enumerate(playerSinguities) if singuityClusterIds[index] != -1]:
            singuityClusters[singuityClusterIds[index]].append(singuity)

        def getClusterForce(singuities: List[Singuity]) -> float:
            clusterStd = arrays.mad([s.position for s in singuities], axis=0, minMad=PhysicsEstimator.getMinimumStd(len(singuities)))
            averageHealth = np.mean([s.healthPoints for s in singuities]).item()
            return PhysicsEstimator.getClusterForce(len(singuities), np.linalg.norm(clusterStd), averageHealth, 0)

        clusterForces = [getClusterForce(cluster) for cluster in singuityClusters]

        if len(clusterForces) == 0:
            singuityStd, singuityCenter = arrays.mad(singuityPositions, axis=0, returnMedian=True, minMad=PhysicsEstimator.getMinimumStd(len(singuityPositions)))
            singuityAverageHealth = np.mean([s.healthPoints for s in playerSinguities]).item()
            return DiscretePlayer(
                currentPlayerId == playerId, playerId, len(playerSinguities), singuityCenter, np.linalg.norm(singuityStd), singuityAverageHealth, [s.id for s in playerSinguities]
                )

        else:
            keptCluster = singuityClusters[np.argmax(clusterForces)]

            singuityStd, singuityCenter = arrays.mad([s.position for s in keptCluster], axis=0, returnMedian=True, minMad=PhysicsEstimator.getMinimumStd(len(keptCluster)))
            singuityAverageHealth = np.mean([s.healthPoints for s in keptCluster]).item()
            return DiscretePlayer(
                currentPlayerId == playerId, playerId, len(keptCluster), singuityCenter, np.linalg.norm(singuityStd), singuityAverageHealth, [s.id for s in keptCluster]
                )

    def executeMovement(self, move: DiscreteMove, targetPosition: np.ndarray, restrictedDuration: int = None, acceptableSinguityCount: int = None):
        def getLastingMoveDuration(moveDuration: int) -> int:
            minimumBoundDuration = moveDuration if move.minimumMoveTime is None else max(moveDuration, move.minimumMoveTime)
            maximumBoundDuration = minimumBoundDuration if restrictedDuration is None else min(minimumBoundDuration, restrictedDuration)
            return maximumBoundDuration

        originalMovementDuration = PhysicsEstimator.estimateMovementDuration(
            self.singuitiesMeanPosition, targetPosition, clusterStd=self.singuitiesStd, acceptableSinguityCount=None if acceptableSinguityCount is None else (
            self.singuityCount, acceptableSinguityCount)
            )
        constrainedMovementDuration = getLastingMoveDuration(originalMovementDuration)

        if constrainedMovementDuration == 0:
            return 0, self

        targetIsWithinStd = PhysicsEstimator.distance2(self.singuitiesMeanPosition, targetPosition) <= self.singuitiesStd ** 2

        if targetIsWithinStd:
            ratioOfTimeSpentInside = min(constrainedMovementDuration * Singuity.MAXIMUM_SPEED_UNITS_PER_FRAME, self.singuitiesStd * 2) / (self.singuitiesStd * 2)
            newStd = min(PhysicsEstimator.getMinimumStd(self.singuityCount) * ratioOfTimeSpentInside + self.singuitiesStd * (1 - ratioOfTimeSpentInside), self.singuitiesStd)
        else:
            # maximum duration on a 5000 unit map
            newStd = self.singuitiesStd * (1 + constrainedMovementDuration / 1768)

        return constrainedMovementDuration,\
               DiscretePlayer(
                   self.isCurrentPlayer,
                   self.id,
                   self.singuityCount,
                   self.singuitiesMeanPosition + (targetPosition - self.singuitiesMeanPosition) * (constrainedMovementDuration / constrainedMovementDuration),
                   newStd,
                   self.singuitiesAverageHealth
               )

    def fought(self, remainingSinguityCount):
        return DiscretePlayer(self.isCurrentPlayer, self.id, remainingSinguityCount, self.singuitiesMeanPosition, self.singuitiesStd, self.singuitiesAverageHealth)

    def tryClaimFor(self, remainingSinguitiesToClaim):
        return DiscretePlayer(
            self.isCurrentPlayer, self.id, max(self.singuityCount - remainingSinguitiesToClaim, 0), self.singuitiesMeanPosition, self.singuitiesStd, self.singuitiesAverageHealth
            )

    def appendNewSpawned(self, spawners: List[Tuple[np.ndarray, int]], anteMoveFrameCount: int, postMoveFrameCount: int) -> DiscretePlayer:
        spawnerLastClaimedFrames = np.array([s[1] for s in spawners])
        spawnerPositions = np.array([s[0] for s in spawners])

        matureSpawnerDurations = np.maximum(0, postMoveFrameCount - np.maximum(anteMoveFrameCount, spawnerLastClaimedFrames + Spawner.GESTATION_FRAME_LAG))
        newSinguityCounts = np.round(matureSpawnerDurations / Spawner.SPAWN_FRAME_LAG)

        usefulSpawnersMask = newSinguityCounts > 0
        if np.count_nonzero(usefulSpawnersMask) <= 0:
            return self

        spawnerPositions = spawnerPositions[usefulSpawnersMask]
        matureSpawnerDurations = matureSpawnerDurations[usefulSpawnersMask]
        newSinguityCounts = newSinguityCounts[usefulSpawnersMask]

        distanceFromSpawnersToMeanPosition = PhysicsEstimator.distance(self.singuitiesMeanPosition, spawnerPositions)
        durationsToTarget = distanceFromSpawnersToMeanPosition / Singuity.MAXIMUM_SPEED_UNITS_PER_FRAME
        atTargetCounts = np.round(np.maximum((matureSpawnerDurations - durationsToTarget) * PhysicsEstimator.SPAWNER_SPAWN_PER_FRAME, 0))

        atTargetCount = np.sum(atTargetCounts)

        furthestSinguityRatios = np.ones_like(distanceFromSpawnersToMeanPosition)
        matureSpawnerDurationIsLowerThanTargetDurationMask = matureSpawnerDurations < atTargetCounts
        furthestSinguityRatios[matureSpawnerDurationIsLowerThanTargetDurationMask] = matureSpawnerDurations[matureSpawnerDurationIsLowerThanTargetDurationMask] / atTargetCounts[
            matureSpawnerDurationIsLowerThanTargetDurationMask]

        inLineCounts = newSinguityCounts - atTargetCounts
        inLineMeanPositions = spawnerPositions + (spawnerPositions - self.singuitiesMeanPosition) * furthestSinguityRatios[:, np.newaxis] / 2
        inLineStds = distanceFromSpawnersToMeanPosition * furthestSinguityRatios / 4
        lineUsageMask = PhysicsEstimator.distance2(spawnerPositions, self.singuitiesMeanPosition) <= (self.singuitiesStd * 2) ** 2

        means = np.concatenate([[self.singuitiesMeanPosition, self.singuitiesMeanPosition], inLineMeanPositions[lineUsageMask]])
        stds = np.concatenate([[self.singuitiesStd, PhysicsEstimator.getMinimumStd(atTargetCount)], inLineStds[lineUsageMask]])
        counts = np.concatenate([[self.singuityCount, atTargetCount], inLineCounts[lineUsageMask]])

        mean, std, count = arrays.combineMeanStdAndCount(means, stds, counts, minStd=np.ones(2))

        return DiscretePlayer(
            self.isCurrentPlayer, self.id, count, mean, np.mean(std),
            0 if count == 0 else (self.singuitiesAverageHealth * self.singuityCount + Singuity.MAX_HEALTH_POINT * (count - self.singuityCount)) / count
            )


'''
test_read_unzip_docx >> True
test_invalid_docx >> True
test_write_default_docx >> True
test_write_name_docx >> True
test_write_save_file >> True
test_invalid_path >> True
'''

import unittest
import os
import shutil
from pathlib import Path

from docxpy.core.extractor import DocxExtractorWrite, DocxExtractorRead
from docxpy.core.exceptions import ErrorPathToDocx, ErrorDocxFile

# Removed sorting of methods in tests.
unittest.TestLoader.sortTestMethodsUsing = None  # type: ignore

# Constants
DIRECTORY = Path(os.path.split(__file__)[0]) / 'templates'

class DocxExtractorTest(unittest.TestCase):
    """
    Tests to check the correct operation of
    the DocxExtractor class functional.
    """

    @classmethod
    def tearDownClass(cls) -> None:
        super().tearDownClass()
        shutil.rmtree(DIRECTORY)

    def test_read_unzip_docx(self) -> None:
        """Check unzipping of a file with docx extension."""
        path = Path(os.getcwd()) / 'tests'
        file_name = 'file.docx'
        self.assertTrue(
            (path / file_name).is_file(),
            'Error! No file: file.docx!',
        )
        DocxExtractorRead(path, file_name).read()
        self.assertTrue(
            (path / 'templates').is_dir(),
            'Error! The file won\'t unpack! Check DocxExtractorRead',
        )

    def test_invalid_docx(self) -> None:
        """
        Check for invalid data in DocxExtractorRead.
        """
        path = Path(os.getcwd()) / 'tests'
        file_name = 'document.docx'
        self.assertRaises(
            (ErrorPathToDocx, ErrorDocxFile),
            DocxExtractorRead,
            path=path,
            name=file_name,
        )

    def test_write_default_docx(self) -> None:
        """
        Check the creation of a file with
        the name default and the extension docx.
        """

        DocxExtractorWrite(DIRECTORY).write()
        self._check_file((DIRECTORY / 'default.docx').is_file())
        (DIRECTORY / 'default.docx').unlink(missing_ok=True)

    def test_write_name_docx(self) -> None:
        """Check the creation of a file with a given name."""

        name_file_docx = 'test.docx'
        DocxExtractorWrite(DIRECTORY, name_file_docx).write()
        self._check_file((DIRECTORY / name_file_docx).is_file())
        (DIRECTORY / name_file_docx).unlink(missing_ok=True)
        name_file = 'test'
        DocxExtractorWrite(DIRECTORY, name_file).write()
        self._check_file((DIRECTORY / name_file_docx).is_file())
        (DIRECTORY / name_file_docx).unlink(missing_ok=True)

    def test_write_save_file(self) -> None:
        """
        Checking that the file is created
        in the required location when you save it.
        """

        directory = Path(os.path.split(__file__)[0]) / 'save_dir'
        directory.mkdir(parents=True, exist_ok=True)
        DocxExtractorWrite(DIRECTORY, path_save=directory).write()
        self._check_file((directory / 'default.docx').is_file())
        shutil.rmtree(directory)

    def test_invalid_path(self) -> None:
        """Check the handling of an invalid path."""

        invalid_path = './test_dir/test_dir/@$'
        self.assertRaises(
            FileExistsError,
            DocxExtractorWrite,
            path=invalid_path,
        )
        self.assertRaises(
            FileExistsError,
            DocxExtractorWrite,
            path=DIRECTORY,
            path_save=invalid_path,
        )

    def _check_file(self, path_bool: bool) -> None:
        self.assertTrue(
            path_bool,
            'Error! The file is not being created! Check DocxExtractorWrite!',
        )