from rest_framework.decorators import api_view
from rest_framework.response import Response

from frontend.models import PCT, Practice
from api.geojson_serializer import as_geojson
import api.view_utils as utils


@api_view(['GET'])
def org_location(request, format=None):
    org_type = request.GET.get('org_type', '')
    centroids = request.GET.get('centroids', '')
    org_codes = utils.param_to_list(request.GET.get('q', ''))
    if org_type == 'practice':
        result_spec = _get_practices(org_codes, centroids)
    elif org_type == 'ccg':
        result_spec = _get_ccgs(org_codes, centroids)
    else:
        raise ValueError('Unknown org_type: {}'.format(org_type))
    results, geo_field, other_fields = result_spec
    fields = other_fields + (geo_field,)
    return Response(
        as_geojson(results.values(*fields), geometry_field=geo_field)
    )


def _get_practices(org_codes, centroids):
    org_codes = utils.get_practice_ids_from_org(org_codes)
    results = Practice.objects.filter(code__in=org_codes)
    geo_field = 'location'
    other_fields = ('name', 'code', 'setting', 'is_dispensing')
    return results, geo_field, other_fields


def _get_ccgs(org_codes, centroids):
    results = PCT.objects.filter(close_date__isnull=True, org_type='CCG')
    if org_codes:
        results = results.filter(code__in=org_codes)
    geo_field = 'centroid' if centroids else 'boundary'
    other_fields = ('name', 'code', 'ons_code', 'org_type')
    return results, geo_field, other_fields

@@ -1,28 +1,87 @@
+from django.contrib.gis.db.models.aggregates import Union%0A%0A
 from rest_framework.decorato
@@ -186,16 +186,35 @@
 Practice
+, STP, RegionalTeam
 %0Afrom ap
@@ -670,24 +670,208 @@
  centroids)%0A
+    elif org_type == 'stp':%0A        result_spec = _get_stps(org_codes, centroids)%0A    elif org_type == 'regional_team':%0A        result_spec = _get_regional_teams(org_codes, centroids)%0A
     else:%0A  
@@ -1743,28 +1743,707 @@
 ts, geo_field, other_fields%0A
+%0A%0Adef _get_stps(org_codes, centroids):%0A    results = STP.objects.all()%0A    if org_codes:%0A        results = results.filter(ons_code__in=org_codes)%0A    results = results.annotate(boundary=Union('pct__boundary'))%0A    geo_field = 'boundary'%0A    other_fields = ('name', 'ons_code')%0A    return results, geo_field, other_fields%0A%0A%0Adef _get_regional_teams(org_codes, centroids):%0A    results = RegionalTeam.objects.filter(close_date__isnull=True)%0A    if org_codes:%0A        results = results.filter(code__in=org_codes)%0A    results = results.annotate(boundary=Union('pct__boundary'))%0A    geo_field = 'boundary'%0A    other_fields = ('name', 'code')%0A    return results, geo_field, other_fields%0A

Support STPs/Regional Teams in location endpoint