from django.core.management.base import BaseCommand
from corehq.apps.domain.models import Domain
from corehq.apps.locations.models import SQLLocation


EXCLUDE_DOMAINS = (
    "drewpsi",
    "psi",
    "psi-ors",
    "psi-test",
    "psi-test2",
    "psi-test3",
    "psi-unicef",
    "psi-unicef-wb",
)


class Command(BaseCommand):
    help = 'Check for consistency between location types and locations'

    def check_for_duplicate_codes(self, loc_types):
        if not len(loc_types) == len({lt.code for lt in loc_types}):
            self.stdout.write("  Duplicate location_type codes found\n")

    def check_single_parentage(self, loc_types):
        for loc_type in loc_types:
            if len(loc_type.allowed_parents) != 1:
                self.stdout.write(u"  {} doesn't have one parent: ({})\n"
                                  .format(loc_type.name, loc_type.allowed_parents))

    def _find_parent(self, parent_code, loc_types):
        for loc_type in loc_types:
            # Some domains have names and not codes. Fall back to that.
            if loc_type.code == parent_code or loc_type.name == parent_code:
                return loc_type
        self.stdout.write(u"  parent loc type {} not found".format(parent_code))

    def check_parentage(self, loc_types):
        for loc_type in loc_types:
            visited = set()
            def step(lt):
                if lt.code in visited:
                    self.stdout.write(u"  I found a cycle!!! {}\n"
                                      .format(loc_types))
                    return
                elif not lt.allowed_parents[0]:
                    return
                else:
                    visited.add(lt.code)
                    parent = self._find_parent(lt.allowed_parents[0], loc_types)
                    if not parent:
                        return
                    step(parent)
            step(loc_type)

    def check_loc_type(self, loc, loc_type_codes):
        if loc.location_type not in loc_type_codes:
            self.stdout.write(
                u"  Location {} has a nonexistant type of {}\n"
                .format(loc.location_id, loc.location_type)
            )

    def handle(self, *args, **options):
        for domain_obj in Domain.get_all():
            loc_types = domain_obj.location_types
            locs = SQLLocation.objects.filter(domain=domain_obj.name).all()
            if not loc_types and not locs:
                continue
            self.stdout.write(u"[{}] - {}\n".format(domain_obj.name,
                                                    domain_obj._id))
            self.check_for_duplicate_codes(loc_types)
            self.check_single_parentage(loc_types)
            self.check_parentage(loc_types)
            loc_type_codes = {loc_type.code for loc_type in loc_types}
            # Some domains have names and not codes. Fall back to that.
            loc_type_codes.update({loc_type.name for loc_type in loc_types})
            for loc in locs:
                self.check_loc_type(loc, loc_type_codes)

@@ -1933,244 +1933,876 @@
 type
+s
 (self, 
-loc, loc_type_codes):%0A        if loc.location_type not in loc_type_codes:%0A            self.stdout.write(%0A                u%22  Location %7B%7D has a nonexistant type of %7B%7D%5Cn%22%0A                .format(loc.location_id, loc.location_type)%0A
+domain, loc_types):%0A        loc_type_codes = %7Bloc_type.code for loc_type in loc_types%7D%0A        # Some domains have names and not codes. Fall back to that.%0A        loc_type_codes.update(%7Bloc_type.name for loc_type in loc_types%7D)%0A        loc_types = (SQLLocation.objects%0A                     .filter(domain=domain)%0A                     .values_list('location_type', flat=True)%0A                     .order_by('location_type')%0A                     .distinct('location_type'))%0A%0A        for loc_type in loc_types:%0A            if loc_type not in loc_type_codes:%0A                count = (SQLLocation.objects%0A                         .filter(domain=domain, location_type=loc_type)%0A                         .count())%0A                self.stdout.write(%0A                    u%22  loc_type %7B%7D does not exist (%7B%7D instances)%5Cn%22%0A                    .format(loc_type, count)%0A    
     
@@ -3016,19 +3016,22 @@
 j.name).
-all
+exists
 ()%0A     
@@ -3396,284 +3396,44 @@
     
-loc_type_codes = %7Bloc_type.code for loc_type in loc_types%7D%0A            # Some domains have names and not codes. Fall back to that.%0A            loc_type_codes.update(%7Bloc_type.name for loc_type in loc_types%7D)%0A            for loc in locs:%0A                self.check_loc_type(loc
+self.check_loc_types(domain_obj.name
 , lo
@@ -3438,16 +3438,11 @@
 loc_type
-_code
 s)%0A

Optimize loc_types queries