# Copyright 2017 Sidewalk Labs | https://www.apache.org/licenses/LICENSE-2.0

from __future__ import (
    absolute_import, division, print_function, unicode_literals
)

import pandas

from doppelganger import inputs


class Population(object):

    @staticmethod
    def from_csvs(persons_infile, households_infile):
        """Load generated population from file.

        Args:
            persons_infile (unicode): path to persons csv
            households_infile (unicode): path to households csv

        Returns:
            Population: generated population

        """
        generated_people = pandas.read_csv(persons_infile)
        generated_households = pandas.read_csv(households_infile)
        return Population(generated_people, generated_households)

    def __init__(self, generated_people, generated_households):
        self.generated_people = generated_people
        self.generated_households = generated_households

    @staticmethod
    def _extract_person_evidence(allocated_rows, fields, segmenter, household_allocator):
        """Creates a (python) generator for bayesian network evidence for persons that yields
         (serial number, evidence, segment, tract, count)
        """
        for _, row in allocated_rows.iterrows():
            serialno = row[inputs.SERIAL_NUMBER.name]
            evidence = tuple((field, row[field]) for field in fields)
            segment = segmenter(row)
            # Draw repeat information for persons from from the allocator
            count_info = household_allocator.get_counts(serialno)
            for tract, count in count_info:
                yield serialno, evidence, segment, tract, count

    @staticmethod
    def _extract_household_evidence(allocated_rows, fields, segmenter, _):
        """Creates a (python) generator for bayesian network evidence for households that yields
         (serial number, evidence, segment, tract, count)
        """
        for _, row in allocated_rows.iterrows():
            serialno = row[inputs.SERIAL_NUMBER.name]
            evidence = tuple((field, row[field]) for field in fields)
            segment = segmenter(row)
            # Households store their repeat information directly
            yield serialno, evidence, segment, row[inputs.TRACT.name], row['count']

    @staticmethod
    def _generate_from_model(household_allocator, data, model, fields, evidence_fn):
        """Generate the given fields of the given data generated by the
        given model
        """
        results = []
        for serialno, evidence, segment, tract, count in evidence_fn(
                data,
                fields,
                model.segmenter,
                household_allocator
        ):
            generated_rows = model.generate(segment, evidence, count=count)
            for repeat_id, row in enumerate(generated_rows):
                household_id = '{}-{}-{}'.format(tract, serialno, repeat_id)
                results.append((household_id, tract, serialno, repeat_id) + row)

        column_names = [inputs.HOUSEHOLD_ID.name, inputs.TRACT.name, inputs.SERIAL_NUMBER.name,
                        inputs.REPEAT_INDEX.name] + list(model.fields)
        results_dataframe = pandas.DataFrame(results, columns=column_names)
        return results_dataframe

    @staticmethod
    def generate(household_allocator, person_model, household_model):
        """Create all the persons and households for this population

        Args:
            household_allocator (HouseholdAllocator): allocated households
            person_model (BayesianNetworkNodel): optional generative model
            household_model (BayesianNetworkNodel): optional generative model

        Returns: Population from the given model
        """

        persons = Population._generate_from_model(
            household_allocator, household_allocator.allocated_persons,
            person_model, [inputs.AGE.name, inputs.SEX.name], Population._extract_person_evidence
        )
        households = Population._generate_from_model(
            household_allocator, household_allocator.allocated_households,
            household_model, [inputs.NUM_PEOPLE.name], Population._extract_household_evidence
        )
        return Population(persons, households)

    def write(self, persons_outfile, households_outfile):
        """Write population to the given file

        Args:
            outfile (unicode): path to write to

        """
        self.generated_people.to_csv(persons_outfile)
        self.generated_households.to_csv(households_outfile)
