import math
import subprocess
import gzip
import shutil

import geopandas as gpd
from shapely.geometry import Point
import hub.helpers.constants as cte
from lxml import etree

CONFIG_DTD = "http://www.matsim.org/files/dtd/config_v2.dtd"
FACILITIES_DTD = "http://www.matsim.org/files/dtd/facilities_v1.dtd"
POPULATION_DTD = "http://www.matsim.org/files/dtd/population_v5.dtd"


class Matsim:
  def __init__(self, city, output_file_path):
    self._city = city
    self._output_file_path = output_file_path

    self._facilities = {
      'name': self._city.name + ' Facilities',
      'facility': []
    }

  def export(self):
    self._export_facilities()
    self._export_network()
    self._export_population()
    self._export_config()

  def _export_facilities(self):
    buildings_shape_data = {
      'id': [],
      'geometry': []
    }

    facilities_xml = etree.Element('facilities', name=self._facilities['name'])

    for building in self._city.buildings:
      for surface in building.grounds:
        for coord in surface.solid_polygon.coordinates:
          buildings_shape_data['id'].append(f"{building.name}")
          buildings_shape_data['geometry'].append(Point(coord[0], coord[1]))

      facility = {
        'id': building.name,
        'x': str(building.centroid[0]),
        'y': str(building.centroid[1]),
        'activity': []
      }

      if len(building.thermal_zones_from_internal_zones) > 1:
        raise NotImplementedError("multi-zone buildings aren't yet supported")

      building_schedules = []

      capacity = 0
      for thermal_zone in building.thermal_zones_from_internal_zones:
        capacity = thermal_zone.occupancy.occupancy_density * building.floor_area * building.storeys_above_ground
        for schedule in thermal_zone.occupancy.occupancy_schedules:
          building_schedules.append(schedule)

      activity_info = {
        'type': building.function,
        'capacity': math.ceil(capacity),
        'opentime': _convert_schedules(building_schedules)
      }

      facility_xml = etree.SubElement(facilities_xml, 'facility', {
        'id': f"{facility['id']}",
        'x': f"{facility['x']}",
        'y': f"{facility['y']}",
      })

      activity_xml = etree.SubElement(facility_xml, 'activity', {
        'type': f"{activity_info['type']}"
      })

      etree.SubElement(activity_xml, 'capacity', {
        'value': f"{activity_info['capacity']}"
      })

      etree.SubElement(activity_xml, 'opentime', {
        'day': f"{activity_info['opentime'][0]['day']}",
        'start_time': f"{activity_info['opentime'][0]['start_time']}",
        'end_time': f"{activity_info['opentime'][0]['end_time']}"
      })

      facility['activity'].append(activity_info)
      self._facilities['facility'].append(facility)

    gdf = gpd.GeoDataFrame(
      buildings_shape_data,
      crs=self._city.srs_name
    )
    gdf.to_file("input_files/buildings_shapefile.shp")

    # Write xml content to file
    xml_content = etree.tostring(facilities_xml, pretty_print=True, encoding='UTF-8').decode('utf-8')
    output_file = f"{self._output_file_path}/{self._city.name}_facilities.xml"
    _save_xml(output_file, xml_content, f"<!DOCTYPE facilities SYSTEM \"{FACILITIES_DTD}\">\n", True)

  def _export_network(self):
    java_path = "java"
    jar_path = "matsim-network-from-osm.jar"
    command = [
      java_path,
      "-jar", jar_path,
      "input_files/merged-network.osm.pbf",
      "input_files/buildings_shapefile.shp",
      f"{self._output_file_path}/{self._city.name}_network.xml.gz"
    ]
    subprocess.run(command)

  def _export_population(self):
    population = etree.Element("population")
    id = 0

    # Generate work facilities
    work = []
    for facility in self._facilities['facility']:
      if facility['activity'][0]['type'] != cte.RESIDENTIAL:
        work.append({
          'type': facility['activity'][0]['type'],
          'capacity': int(facility['activity'][0]['capacity']),
          'facility': facility['id'],
          'x': facility['x'],
          'y': facility['y'],
          'start_time': '08:00:00',
          'end_time': '18:00:00'
        })

    # Generate the population from residential places first
    current_work = 0
    for facility in self._facilities['facility']:
      if facility['activity'][0]['type'] == cte.RESIDENTIAL:
        max_capacity = int(facility['activity'][0]['capacity'])
        for i in range(max_capacity):
          person = etree.SubElement(population, 'person', {
            'id': str(id),
            'sex': 'm',
            'age': '32',
            'car_avail': 'always',
            'employed': 'yes',
          })
          plan = etree.SubElement(person, 'plan', {'selected': 'yes'})

          # Residential activity
          etree.SubElement(plan, 'act', {
            'type': f"{facility['activity'][0]['type']}",
            'facility': f"{facility['id']}",
            'x': f"{facility['x']}",
            'y': f"{facility['y']}",
            'end_time': '7:30:00'
          })

          # Leg to work
          etree.SubElement(plan, 'leg', {'mode': 'car'})

          # Work activity
          etree.SubElement(plan, 'act', {
            'type': f"{work[current_work]['type']}",
            'facility': f"{work[current_work]['facility']}",
            'x': f"{work[current_work]['x']}",
            'y': f"{work[current_work]['y']}",
            'start_time': f"{work[current_work]['start_time']}",
            'end_time': f"{work[current_work]['end_time']}",
          })

          # Leg to home
          etree.SubElement(plan, 'leg', {'mode': 'car'})

          # Residential activity (return)
          etree.SubElement(plan, 'act', {
            'type': f"{facility['activity'][0]['type']}",
            'facility': f"{facility['id']}",
            'x': f"{facility['x']}",
            'y': f"{facility['y']}",
          })

          work[current_work]['capacity'] -= 1
          if work[current_work]['capacity'] == 0:
            current_work += 1

          id += 1

    # Write xml content to file
    xml_content = etree.tostring(population, pretty_print=True, encoding='UTF-8').decode('utf-8')
    output_file = f"{self._output_file_path}/{self._city.name}_population.xml"
    _save_xml(output_file, xml_content, f"<!DOCTYPE population SYSTEM \"{POPULATION_DTD}\">\n", True)

  def _export_config(self):
    root = etree.Element('config')

    # ======== NETWORK ========= #
    network_path_module = etree.SubElement(root, 'module', {
      'name': 'network'
    })
    _add_param(network_path_module, 'inputNetworkFile', f"{self._city.name}_network.xml.gz")

    # ======== POPULATION ========= #
    population_path_module = etree.SubElement(root, 'module', {
      'name': 'plans'
    })
    _add_param(population_path_module, 'inputPlansFile', f"{self._city.name}_population.xml.gz")

    # ======== FACILITIES ========= #
    facilities_path_module = etree.SubElement(root, 'module', {
      'name': 'facilities'
    })
    _add_param(facilities_path_module, 'inputFacilitiesFile', f"{self._city.name}_facilities.xml.gz")
    _add_param(facilities_path_module, 'facilitiesSource', 'fromFile')

    # ======== CONTROLER ========= #
    controler_module = etree.SubElement(root, 'module', {
      'name': 'controler'
    })
    controler_params = [
      ('writeEventsInterval', '1000'),
      ('writePlansInterval', '1000'),
      ('eventsFileFormat', 'xml'),
      ('outputDirectory', f"{self._output_file_path}/Montreal"),
      ('firstIteration', '0'),
      ('lastIteration', '10'),
      ('mobsim', 'qsim'),
    ]
    for param in controler_params:
      _add_param(controler_module, param[0], param[1])

    # ======== QSIM ========= #
    qsim_module = etree.SubElement(root, 'module', {
      'name': 'qsim'
    })
    qsim_params = [
      ('startTime', '00:00:00'),
      ('endTime', '30:00:00'),
      ('flowCapacityFactor', '1.00'),
      ('storageCapacityFactor', '1.00'),
      ('numberOfThreads', '1'),
      ('snapshotperiod', '00:00:01'),
      ('removeStuckVehicles', 'false'),
      ('stuckTime', '3600.0'),
      ('timeStepSize', '00:00:01'),
      ('trafficDynamics', 'queue'),
    ]
    for param in qsim_params:
      _add_param(qsim_module, param[0], param[1])

    # ======== SCORING ========= #
    score_module = etree.SubElement(root, 'module', {
      'name': 'planCalcScore'
    })
    _add_param(score_module, 'BrainExpBeta', '1.0')
    _add_param(score_module, 'learningRate', '1.0')

    scoring_paramset = etree.SubElement(score_module, 'parameterset', {
      'type': 'scoringParameters'
    })
    scoring_paramset_params = [
      ('earlyDeparture', '0.0'),
      ('lateArrival', '0.0'),
      ('marginalUtilityOfMoney', '0.062'),
      ('performing', '0.96'),
      ('utilityOfLineSwitch', '0.0'),
      ('waitingPt', '-0.18'),
    ]
    for param in scoring_paramset_params:
      _add_param(scoring_paramset, param[0], param[1])

    mode_paramsets = [
      ("car",
       [
         ("marginalUtilityOfTraveling_util_hr", "0.0"),
         ("constant", "-0.562"),
         ("monetaryDistanceRate", "-0.0004")
       ]),
      ("walk",
       [
         ("marginalUtilityOfTraveling_util_hr", "-1.14"),
         ("constant", "0.0"),
         ("marginalUtilityOfDistance_util_m", "0.0")
       ])
    ]
    for mode, parameters in mode_paramsets:
      _add_parameterset(scoring_paramset, "modeParams", [("mode", mode)] + parameters)

    activity_paramsets = [
      ('residential',[
        ('priority', '1'),
        ('typicalDuration', '13:00:00'),
        ('minimalDuration', '01:00:00'),
      ]),
      ('medium office',[
        ('priority', '1'),
        ('typicalDuration', '09:00:00'),
        ('minimalDuration', '08:00:00'),
        ('openingTime', '08:00:00'),
        ('earliestEndTime', '17:00:00'),
        ('latestStartTime', '09:00:00'),
        ('closingTime', '18:00:00'),
      ]),
      ('warehouse',[
        ('priority', '1'),
        ('typicalDuration', '09:00:00'),
        ('minimalDuration', '08:00:00'),
        ('openingTime', '08:00:00'),
        ('earliestEndTime', '17:00:00'),
        ('latestStartTime', '09:00:00'),
        ('closingTime', '18:00:00'),
      ]),
      ('stand alone retail',[
        ('priority', '1'),
        ('typicalDuration', '09:00:00'),
        ('minimalDuration', '08:00:00'),
        ('openingTime', '08:00:00'),
        ('earliestEndTime', '17:00:00'),
        ('latestStartTime', '09:00:00'),
        ('closingTime', '18:00:00'),
      ]),
    ]
    for activity_type, parameters in activity_paramsets:
      _add_parameterset(scoring_paramset, "activityParams", [("activityType", activity_type)] + parameters)

    # ======== STRATEGY ========= #
    strategy_module = etree.SubElement(root, 'module', {
      'name': 'strategy'
    })
    _add_param(strategy_module, 'maxAgentPlanMemorySize', '6')

    strategy_paramsets = [
      ("ChangeExpBeta",[
         ("weight", "0.7"),
       ]),
      ("ReRoute",[
         ("disableAfterIteration", "2900"),
         ("weight", "0.01"),
       ]),
      ("SubtourModeChoice", [
        ("disableAfterIteration", "2900"),
        ("weight", "0.01"),
      ]),
      ("TimeAllocationMutator", [
        ("disableAfterIteration", "2900"),
        ("weight", "0.01"),
      ]),
    ]
    for name, parameters in strategy_paramsets:
      _add_parameterset(strategy_module, "strategysettings", [("strategyName", name)] + parameters)

    # ======== SUBTOUR MODE CHOICE ========= #
    subtour_module = etree.SubElement(root, 'module', {
      'name': 'subtourModeChoice'
    })
    # Defines the chain-based modes, seperated by commas
    _add_param(subtour_module, 'chainBasedModes', 'car')
    # Defines whether car availability must be considered or not. An agent has no car only if it has no license, or never access to a car
    _add_param(subtour_module, 'considerCarAvailability', 'true')
    # Defines all the modes available, including chain-based modes, seperated by commas
    _add_param(subtour_module, 'modes', 'car,walk')

    # Write xml content to file
    xml_content = etree.tostring(root, pretty_print=True, encoding='UTF-8').decode('utf-8')
    output_file = f"{self._output_file_path}/{self._city.name}_config.xml"
    _save_xml(output_file, xml_content, f"<!DOCTYPE config SYSTEM \"{CONFIG_DTD}\">\n")


def _add_param(parent, name, value):
  etree.SubElement(parent, "param", {
    'name': name,
    'value': value
  })


def _add_parameterset(parent, type, parameters):
  parameterset = etree.SubElement(parent, "parameterset", {
    'type': type
  })
  for name, value in parameters:
    _add_param(parameterset, name, value)


def _convert_schedules(building_schedules):
  converted_schedules = []
  opening_hour = 0
  closing_hour = 0
  schedule = building_schedules[0]

  for i, value in enumerate(schedule.values):
    if value > 0:
      opening_hour = i
      break

  for i, value in reversed(list(enumerate(schedule.values))):
    if value > 0:
      closing_hour = i
      break

  for day in schedule.day_types:
    if day[0:3] != 'hol':
      converted_schedules.append({
        'day': day[0:3],
        'start_time': f"{opening_hour:02}:00:00",
        'end_time': f"{closing_hour:02}:00:00"
      })

  return converted_schedules


def _save_xml(output_file, xml_content, xml_dtd=None, zipped=None):
  if zipped is None:
    zipped = False

  if xml_dtd is None:
    xml_dtd = ''

  with open(output_file, 'w') as file:
    file.write("<?xml version=\"1.0\" encoding=\"utf-8\"?>\n")
    file.write(xml_dtd)
    file.write(xml_content)

  if zipped:
    with open(output_file, 'rb') as f_in:
      with gzip.open(output_file + '.gz', 'wb') as f_out:
        shutil.copyfileobj(f_in, f_out)