<commit_before>"""
Usage:

    ./manage.py import_transxchange EA.zip [EM.zip etc]
"""

import logging
import warnings
import os
import csv
import yaml
import zipfile
import xml.etree.cElementTree as ET
from datetime import timedelta
from django.conf import settings
from django.core.management.base import BaseCommand
from django.contrib.gis.geos import LineString, MultiLineString
from django.db import transaction
from django.db.models import F
from django.utils import timezone
from busstops.models import (Operator, Service, DataSource, StopPoint, StopUsage, ServiceCode)
from ...models import Route, Calendar, CalendarDate, Trip, StopTime, Note
from timetables.txc import TransXChange, Grouping, sanitize_description_part


logger = logging.getLogger(__name__)


NS = {'txc': 'http://www.transxchange.org.uk/'}
BANK_HOLIDAYS = {
    'ChristmasDay': '2019-12-25',
    'BoxingDay':  '2019-12-26',
    'NewYearsDay':  '2020-01-01',
    'AllBankHolidays': '2020-01-01',
    'GoodFriday': '2020-04-10',
    'EasterMonday': '2020-04-13',
}


def sanitize_description(name):
    """
    Given an oddly formatted description from the North East,
    like 'Bus Station bay 5,Blyth - Grange Road turning circle,Widdrington Station',
    returns a shorter, more normal version like
    'Blyth - Widdrington Station'
    """

    parts = [sanitize_description_part(part) for part in name.split(' - ')]
    return ' - '.join(parts)


def get_line_name_and_brand(service_element):
    line_name = service_element.find('txc:Lines', NS)[0][0].text
    if '|' in line_name:
        line_name, line_brand = line_name.split('|', 1)
    else:
        line_brand = ''
    return line_name, line_brand


def get_service_code(filename):
    """
    Given a filename like 'ea_21-45A-_-y08-1.xml',
    returns a service_code like 'ea_21-45A-_-y08'
    """
    parts = filename.split('-')  # ['ea_21', '3', '_', '1']
    if len(parts) == 5:
        net = parts[0].split('_')[0]
        if len(net) <= 3 and net.islower():
            return '-'.join(parts[:-1])


def get_operator_code(operator_element, element_name):
    element = operator_element.find('txc:{}'.format(element_name), NS)
    if element is not None:
        return element.text


def get_operator_name(operator_element):
    "Given an Operator element, returns the operator name or None"

    for element_name in ('TradingName', 'OperatorNameOnLicence', 'OperatorShortName'):
        name = get_operator_code(operator_element, element_name)
        if name:
            return name.replace('&amp;', '&')


def get_operator_by(scheme, code):
    if code:
        return Operator.objects.filter(operatorcode__code=code, operatorcode__source__name=scheme).first()


def line_string_from_journeypattern(journeypattern, stops):
    points = []
    stop = stops.get(journeypattern.sections[0].timinglinks[0].origin.stop.atco_code)
    if stop and stop.latlong:
        points.append(stop.latlong)
    for timinglink in journeypattern.get_timinglinks():
        stop = stops.get(timinglink.destination.stop.atco_code)
        if stop and stop.latlong:
            points.append(stop.latlong)
    try:
        linestring = LineString(points)
        return linestring
    except ValueError:
        pass


class Command(BaseCommand):
    @staticmethod
    def add_arguments(parser):
        parser.add_argument('archives', nargs=1, type=str)
        parser.add_argument('files', nargs='*', type=str)

    def handle(self, *args, **options):
        self.calendar_cache = {}
        self.undefined_holidays = set()
        self.notes = {}
        for archive_name in options['archives']:
            self.handle_archive(archive_name, options['files'])
        if self.undefined_holidays:
            print(self.undefined_holidays)

    def set_region(self, archive_name):
        self.region_id, _ = os.path.splitext(os.path.basename(archive_name))

        self.source, created = DataSource.objects.get_or_create(name=self.region_id)

        if self.region_id == 'NCSD':
            self.region_id = 'GB'

    def get_operator(self, operator_element):
        "Given an Operator element, returns an operator code for an operator that exists."

        for tag_name, scheme in (
            ('NationalOperatorCode', 'National Operator Codes'),
            ('LicenceNumber', 'Licence'),
        ):
            operator_code = get_operator_code(operator_element, tag_name)
            if operator_code:
                operator = get_operator_by(scheme, operator_code)
                if operator:
                    return operator

        # Get by regional operator code
        operator_code = get_operator_code(operator_element, 'OperatorCode')
        if operator_code:
            operator = get_operator_by(self.region_id, operator_code)
            if not operator:
                operator = get_operator_by('National Operator Codes', operator_code)
            if operator:
                return operator

        if get_operator_name(operator_element) not in {'Replacement Service', 'UNKWN'}:
            warnings.warn('Operator not found:\n{}'.format(ET.tostring(operator_element).decode()))

    def get_operators(self, transxchange):
        operators = transxchange.operators
        if len(operators) > 1:
            journey_operators = {journey.operator for journey in transxchange.journeys}
            journey_operators.add(transxchange.operator)
            operators = [operator for operator in operators if operator.get('id') in journey_operators]
        operators = (self.get_operator(operator) for operator in operators)
        return [operator for operator in operators if operator]

    def set_service_descriptions(self, archive):
        # the NCSD has service descriptions in a separate file:
        self.service_descriptions = {}
        if 'IncludedServices.csv' in archive.namelist():
            with archive.open('IncludedServices.csv') as csv_file:
                reader = csv.DictReader(line.decode('utf-8') for line in csv_file)
                # e.g. {'NATX323': 'Cardiff - Liverpool'}
                for row in reader:
                    key = f"{row['Operator']}{row['LineName']}{row['Dir']}"
                    self.service_descriptions[key] = row['Description']

    def get_service_descriptions(self, filename):<commit_msg>Isle of Man<commit_after>"""
Usage:

    ./manage.py import_transxchange EA.zip [EM.zip etc]
"""

import logging
import warnings
import os
import csv
import yaml
import zipfile
import xml.etree.cElementTree as ET
from datetime import timedelta
from django.conf import settings
from django.core.management.base import BaseCommand
from django.contrib.gis.geos import LineString, MultiLineString
from django.db import transaction
from django.db.models import F
from django.utils import timezone
from busstops.models import (Operator, Service, DataSource, StopPoint, StopUsage, ServiceCode)
from ...models import Route, Calendar, CalendarDate, Trip, StopTime, Note
from timetables.txc import TransXChange, Grouping, sanitize_description_part


logger = logging.getLogger(__name__)


NS = {'txc': 'http://www.transxchange.org.uk/'}
BANK_HOLIDAYS = {
    'ChristmasDay': '2019-12-25',
    'BoxingDay':  '2019-12-26',
    'NewYearsDay':  '2020-01-01',
    'AllBankHolidays': '2020-01-01',
    'GoodFriday': '2020-04-10',
    'EasterMonday': '2020-04-13',
}


def sanitize_description(name):
    """
    Given an oddly formatted description from the North East,
    like 'Bus Station bay 5,Blyth - Grange Road turning circle,Widdrington Station',
    returns a shorter, more normal version like
    'Blyth - Widdrington Station'
    """

    parts = [sanitize_description_part(part) for part in name.split(' - ')]
    return ' - '.join(parts)


def get_line_name_and_brand(service_element):
    line_name = service_element.find('txc:Lines', NS)[0][0].text
    if '|' in line_name:
        line_name, line_brand = line_name.split('|', 1)
    else:
        line_brand = ''
    return line_name, line_brand


def get_service_code(filename):
    """
    Given a filename like 'ea_21-45A-_-y08-1.xml',
    returns a service_code like 'ea_21-45A-_-y08'
    """
    parts = filename.split('-')  # ['ea_21', '3', '_', '1']
    if len(parts) == 5:
        net = parts[0].split('_')[0]
        if len(net) <= 3 and net.islower():
            return '-'.join(parts[:-1])


def get_operator_code(operator_element, element_name):
    element = operator_element.find('txc:{}'.format(element_name), NS)
    if element is not None:
        return element.text


def get_operator_name(operator_element):
    "Given an Operator element, returns the operator name or None"

    for element_name in ('TradingName', 'OperatorNameOnLicence', 'OperatorShortName'):
        name = get_operator_code(operator_element, element_name)
        if name:
            return name.replace('&amp;', '&')


def get_operator_by(scheme, code):
    if code:
        return Operator.objects.filter(operatorcode__code=code, operatorcode__source__name=scheme).first()


def line_string_from_journeypattern(journeypattern, stops):
    points = []
    stop = stops.get(journeypattern.sections[0].timinglinks[0].origin.stop.atco_code)
    if stop and stop.latlong:
        points.append(stop.latlong)
    for timinglink in journeypattern.get_timinglinks():
        stop = stops.get(timinglink.destination.stop.atco_code)
        if stop and stop.latlong:
            points.append(stop.latlong)
    try:
        linestring = LineString(points)
        return linestring
    except ValueError:
        pass


class Command(BaseCommand):
    @staticmethod
    def add_arguments(parser):
        parser.add_argument('archives', nargs=1, type=str)
        parser.add_argument('files', nargs='*', type=str)

    def handle(self, *args, **options):
        self.calendar_cache = {}
        self.undefined_holidays = set()
        self.notes = {}
        for archive_name in options['archives']:
            self.handle_archive(archive_name, options['files'])
        if self.undefined_holidays:
            print(self.undefined_holidays)

    def set_region(self, archive_name):
        self.region_id, _ = os.path.splitext(os.path.basename(archive_name))

        self.source, created = DataSource.objects.get_or_create(name=self.region_id)

        if self.region_id == 'NCSD':
            self.region_id = 'GB'
        elif self.region_id == 'IOM':
            self.region_id = 'IM'

    def get_operator(self, operator_element):
        "Given an Operator element, returns an operator code for an operator that exists."

        for tag_name, scheme in (
            ('NationalOperatorCode', 'National Operator Codes'),
            ('LicenceNumber', 'Licence'),
        ):
            operator_code = get_operator_code(operator_element, tag_name)
            if operator_code:
                operator = get_operator_by(scheme, operator_code)
                if operator:
                    return operator

        # Get by regional operator code
        operator_code = get_operator_code(operator_element, 'OperatorCode')
        if operator_code:
            operator = get_operator_by(self.region_id, operator_code)
            if not operator:
                operator = get_operator_by('National Operator Codes', operator_code)
            if operator:
                return operator

        name = get_operator_name(operator_element)

        if name == '<NAME>':
            return Operator.objects.get(name=name)

        if name not in {'Replacement Service', 'UNKWN'}:
            warnings.warn('Operator not found:\n{}'.format(ET.tostring(operator_element).decode()))

    def get_operators(self, transxchange):
        operators = transxchange.operators
        if len(operators) > 1:
            journey_operators = {journey.operator for journey in transxchange.journeys}
            journey_operators.add(transxchange.operator)
            operators = [operator for operator in operators if operator.get('id') in journey_operators]
        operators = (self.get_operator(operator) for operator in operators)
        return [operator for operator in operators if operator]

    def set_service_descriptions(self, archive):
        # the NCSD has service descriptions in a separate file:
        self.service_descriptions = {}
        if 'IncludedServices.csv' in archive.namelist():
            with archive.open('IncludedServices.csv') as csv_file:
                reader = csv.DictReader(line.decode('utf-8') for line in csv_file)
                # e.g. {'NATX323': 'Cardiff - Liverpool'}
                for row in reader:
                    key = f"{row['Operator']}{row['LineName']}{row['Dir']}"
                    self.service_descriptions[key] = row['Description']

    def get_service_descriptions(self, filename):
{'fraction_non_alphanumeric': 0.1084986720824871, 'fraction_numerical': 0.012810498359631307, 'mean_word_length': 3.0172576090367116, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 4}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '524447', 'n_tokens_mistral': 3938, 'n_tokens_neox': 3722, 'n_words': 913}