import csv

from django.core.management.base import BaseCommand

from specimens.models import TaxonRank, Taxon, TaxonStatus, SPECIES_RANK_NAME, SUBGENUS_RANK_NAME

MODELS_TO_TRUNCATE = [Taxon, TaxonRank, TaxonStatus]


def create_initial_ranks():
    TaxonRank.objects.bulk_create([
        TaxonRank(name='Kingdom'),
        TaxonRank(name='Phylum'),
        TaxonRank(name='Class'),
        TaxonRank(name='Order'),
        TaxonRank(name='Family'),
        TaxonRank(name='Genus'),
        TaxonRank(name=SUBGENUS_RANK_NAME),
        TaxonRank(name=SPECIES_RANK_NAME)
    ])

class Command(BaseCommand):
    help = 'Import taxonomy from a CSV file.'

    def add_arguments(self, parser):
        parser.add_argument('csv_file')

        parser.add_argument(
            '--truncate',
            action='store_true',
            dest='truncate',
            default=False,
            help='Truncate all tables prior to import',
        )

    def handle(self, *args, **options):
        self.stdout.write('Importing data from file...')

        with open(options['csv_file']) as csv_file:
            if options['truncate']:
                for model in MODELS_TO_TRUNCATE:
                    self.stdout.write('Truncate model {name} ...'.format(name=model.__name__), ending='')
                    model.objects.all().delete()
                    self.stdout.write(self.style.SUCCESS('OK'))

            self.stdout.write('Creating initial ranks...')
            create_initial_ranks()

            for i, row in enumerate(csv.DictReader(csv_file, delimiter=',')):
                self.stdout.write('Processing row #{i}...'.format(i=i), ending='')

                species_status, _ = TaxonStatus.objects.get_or_create(name=row['Status'])

                # Starting from the higher ranks
                kingdom, _ = Taxon.objects.get_or_create(name=row['Kingdom'].strip(),
                                                         rank=TaxonRank.objects.get(name="Kingdom"))
                phylum, _ = Taxon.objects.get_or_create(name=row['Phylum'].strip(),
                                                        rank=TaxonRank.objects.get(name="Phylum"),
                                                        parent=kingdom)
                class_taxon, _ = Taxon.objects.get_or_create(name=row['Class'].strip(),
                                                          rank=TaxonRank.objects.get(name="Class"),
                                                          parent=phylum)
                order_taxon, _ = Taxon.objects.get_or_create(name=row['Order'].strip(),
                                                             rank=TaxonRank.objects.get(name="Order"),
                                                             parent=class_taxon)
                family, _ = Taxon.objects.get_or_create(name=row['Family'].strip(),
                                                        rank=TaxonRank.objects.get(name="Family"),
                                                        parent=order_taxon)
                genus, _ = Taxon.objects.get_or_create(name=row['Genus'].strip(),
                                                        rank=TaxonRank.objects.get(name="Genus"),
                                                        parent=family)
                # Subgenus rank is optional
                subgenus_source = row['Subgenus'].strip()
                if subgenus_source:
                    subgenus, _ = Taxon.objects.get_or_create(name=subgenus_source,
                                                              rank=TaxonRank.objects.get(name=SUBGENUS_RANK_NAME),
                                                              parent=genus)

                species_parent = subgenus if subgenus_source else genus
                species, _ = Taxon.objects.get_or_create(name=row['Species'].strip(),
                                                         rank=TaxonRank.objects.get(name=SPECIES_RANK_NAME),
                                                         parent=species_parent,
                                                         status=species_status,
                                                         aphia_id=row['Aphia_ID'].strip(),
                                                         authority=row['Authority'].strip())


                self.stdout.write(self.style.SUCCESS('OK'))