from datetime import date, time

from factory import (DjangoModelFactory, Sequence, SubFactory, LazyFunction,
                     post_generation)
from factory.fuzzy import FuzzyText, FuzzyDate, FuzzyInteger, FuzzyFloat
from factory.django import FileField

from ..models import (Flaw, Experiment, ProtocolAttachment, TreatmentType,
                      Treatment, TreatmentReplicate, AliveDeadCount)
from ccdb.misc.tests.factories import ContainerFactory
from ccdb.locations.tests.factories import StudyLocationFactory
from ccdb.species.tests.factories import SpeciesFactory


class FlawFactory(DjangoModelFactory):
    class Meta:
        model = Flaw

    name = Sequence(lambda n: 'flaw{}'.format(n))
    description = FuzzyText(length=255)
    sort_order = Sequence(lambda n: n)


class ExperimentFactory(DjangoModelFactory):
    class Meta:
        model = Experiment

    name = Sequence(lambda n: 'experiment{}'.format(n))
    code = Sequence(lambda n: 'e{}'.format(n))
    description = FuzzyText(length=255)
    flaw = SubFactory(FlawFactory)
    sort_order = Sequence(lambda n: n)

    @post_generation
    def collections(self, create, extracted, **kwargs):  # pragma: no cover
        if not create:
            return
        if extracted:
            for collection in extracted:
                self.groups.add(collection)


class ProtocolAttachmentFactory(DjangoModelFactory):
    class Meta:
        model = ProtocolAttachment

    experiment = SubFactory(ExperimentFactory)
    protocol = FileField()


class TreatmentTypeFactory(DjangoModelFactory):
    class Meta:
        model = TreatmentType

    experiment = SubFactory(ExperimentFactory)
    name = Sequence(lambda n: 'treatment_type{}'.format(n))
    code = Sequence(lambda n: 'tt{}'.format(n))
    treatment_type = FuzzyText(length=50)
    placement = FuzzyText(length=25)
    description = FuzzyText(length=255)
    sort_order = Sequence(lambda n: n)


class TreatmentFactory(DjangoModelFactory):
    class Meta:
        model = Treatment

    treatment_type = SubFactory(TreatmentTypeFactory)
    container = SubFactory(ContainerFactory)
    study_location = SubFactory(StudyLocationFactory)
    species = SubFactory(SpeciesFactory)
    sex = FuzzyText(length=25)
    flaw = SubFactory(FlawFactory)


class TreatmentReplicateFactory(DjangoModelFactory):
    class Meta:
        model = TreatmentReplicate

    treatment = SubFactory(TreatmentFactory)
    name = Sequence(lambda n: 'treatment_replicate{}'.format(n))
    setup_date = FuzzyDate(date(2012, 1, 1))
    setup_time = LazyFunction(time)
    setup_sample_size = FuzzyInteger(0)
    mass_g = FuzzyFloat(0.0)
    flaw = SubFactory(FlawFactory)


class AliveDeadCountFactory(DjangoModelFactory):
    class Meta:
        model = AliveDeadCount

    treatment_replicate = SubFactory(TreatmentReplicateFactory)
    status_date = FuzzyDate(date(2012, 1, 1))
    status_time = LazyFunction(time)
    count_alive = FuzzyInteger(0)
    count_dead = FuzzyInteger(0)
    flaw = SubFactory(FlawFactory)