Treatment n<->1: manage saving

4 files changed, 264 insertions, 70 deletions

diff --git a/archaeological_finds/forms.py b/archaeological_finds/forms.py
index a221247e0..958b1896c 100644
--- a/archaeological_finds/forms.py
+++ b/archaeological_finds/forms.py
@@ -326,90 +326,93 @@ class ResultingFindForm(CustomForm, ManageOldType):
     form_admin_name = _(u"Treatment n-1 - 030 - General")
     form_slug = "treatmentn1-030-general"
 
-    base_models = ['object_type', 'material_type', 'communicabilitie']
     associated_models = {
-        'material_type': models.MaterialType,
-        'object_type': models.ObjectType,
-        'communicabilitie': models.CommunicabilityType,
-        'material_type_quality': models.MaterialTypeQualityType,
-        'object_type_quality': models.ObjectTypeQualityType,
-        'checked_type': models.CheckedType,
+        'resulting_material_type': models.MaterialType,
+        'resulting_object_type': models.ObjectType,
+        'resulting_communicabilitie': models.CommunicabilityType,
+        'resulting_material_type_quality': models.MaterialTypeQualityType,
+        'resulting_object_type_quality': models.ObjectTypeQualityType,
+        'resulting_checked_type': models.CheckedType,
     }
     HEADERS = {}
-    HEADERS['label'] = FormHeader(_(u"Identification"))
+    HEADERS['resulting_label'] = FormHeader(_(u"Identification"))
 
-    label = forms.CharField(
+    resulting_label = forms.CharField(
         label=_(u"Free ID"),
         validators=[validators.MaxLengthValidator(60)])
-    denomination = forms.CharField(label=_(u"Denomination"), required=False)
-    previous_id = forms.CharField(label=_("Previous ID"), required=False)
-    museum_id = forms.CharField(label=_(u"Museum ID"), required=False)
-    seal_number = forms.CharField(label=_(u"Seal number"), required=False)
-    mark = forms.CharField(label=_(u"Mark"), required=False)
+    resulting_denomination = forms.CharField(label=_(u"Denomination"),
+                                             required=False)
 
-    HEADERS['description'] = FormHeader(_(u"Description"))
-    description = forms.CharField(label=_(u"Description"),
-                                  widget=forms.Textarea, required=False)
-    is_complete = forms.NullBooleanField(label=_(u"Is complete?"),
-                                         required=False)
-    material_type = widgets.Select2MultipleField(
+    HEADERS['resulting_description'] = FormHeader(_(u"Description"))
+    resulting_description = forms.CharField(
+        label=_(u"Description"), widget=forms.Textarea, required=False)
+    resulting_is_complete = forms.NullBooleanField(
+        label=_(u"Is complete?"), required=False)
+    resulting_material_type = widgets.Select2MultipleField(
         label=_(u"Material types"), required=False
     )
-    material_type_quality = forms.ChoiceField(
+    resulting_material_type_quality = forms.ChoiceField(
         label=_(u"Material type quality"), required=False, choices=[])
-    object_type = widgets.Select2MultipleField(
+    resulting_object_type = widgets.Select2MultipleField(
         label=_(u"Object types"), required=False,
     )
-    object_type_quality = forms.ChoiceField(
+    resulting_object_type_quality = forms.ChoiceField(
         label=_(u"Object type quality"), required=False, choices=[])
-    find_number = forms.IntegerField(label=_(u"Find number"), required=False)
-    min_number_of_individuals = forms.IntegerField(
+    resulting_find_number = forms.IntegerField(
+        label=_(u"Find number"), required=False)
+    resulting_min_number_of_individuals = forms.IntegerField(
         label=_(u"Minimum number of individuals (MNI)"), required=False)
 
-    decoration = forms.CharField(label=_(u"Decoration"), widget=forms.Textarea,
-                                 required=False)
-    inscription = forms.CharField(label=_(u"Inscription"),
-                                  widget=forms.Textarea, required=False)
-    manufacturing_place = forms.CharField(
+    resulting_decoration = forms.CharField(
+        label=_(u"Decoration"), widget=forms.Textarea, required=False)
+    resulting_inscription = forms.CharField(
+        label=_(u"Inscription"), widget=forms.Textarea, required=False)
+    resulting_manufacturing_place = forms.CharField(
         label=_(u"Manufacturing place"), required=False)
-    communicabilitie = widgets.Select2MultipleField(
+    resulting_communicabilitie = widgets.Select2MultipleField(
         label=_(u"Communicability"), required=False
     )
-    comment = forms.CharField(label=_(u"Comment"), required=False,
+    resulting_comment = forms.CharField(label=_(u"Comment"), required=False,
                               widget=forms.Textarea)
-    dating_comment = forms.CharField(
+    resulting_dating_comment = forms.CharField(
         label=_(u"Comment on dating"), required=False, widget=forms.Textarea)
 
-    HEADERS['length'] = FormHeader(_(u"Dimensions"))
-    length = FloatField(label=_(u"Length (cm)"), required=False)
-    width = FloatField(label=_(u"Width (cm)"), required=False)
-    height = FloatField(label=_(u"Height (cm)"), required=False)
-    diameter = FloatField(label=_(u"Diameter (cm)"), required=False)
-    thickness = FloatField(label=_(u"Thickness (cm)"), required=False)
-    volume = FloatField(label=_(u"Volume (l)"), required=False)
-    weight = FloatField(label=_(u"Weight (g)"), required=False)
-    clutter_long_side = FloatField(
+    HEADERS['resulting_length'] = FormHeader(_(u"Dimensions"))
+    resulting_length = FloatField(label=_(u"Length (cm)"), required=False)
+    resulting_width = FloatField(label=_(u"Width (cm)"), required=False)
+    resulting_height = FloatField(label=_(u"Height (cm)"), required=False)
+    resulting_diameter = FloatField(label=_(u"Diameter (cm)"), required=False)
+    resulting_thickness = FloatField(label=_(u"Thickness (cm)"), required=False)
+    resulting_volume = FloatField(label=_(u"Volume (l)"), required=False)
+    resulting_weight = FloatField(label=_(u"Weight (g)"), required=False)
+    resulting_clutter_long_side = FloatField(
         label=_(u"Clutter long side (cm)"), required=False)
-    clutter_short_side = FloatField(
+    resulting_clutter_short_side = FloatField(
         label=_(u"Clutter short side (cm)"), required=False)
-    clutter_height = FloatField(
+    resulting_clutter_height = FloatField(
         label=_(u"Clutter height (cm)"), required=False)
-    dimensions_comment = forms.CharField(
+    resulting_dimensions_comment = forms.CharField(
         label=_(u"Dimensions comment"), required=False, widget=forms.Textarea)
 
-    HEADERS['checked_type'] = FormHeader(_(u"Sheet"))
-    checked_type = forms.ChoiceField(label=_(u"Check"), required=False)
-    check_date = forms.DateField(
+    HEADERS['resulting_checked_type'] = FormHeader(_(u"Sheet"))
+    resulting_checked_type = forms.ChoiceField(label=_(u"Check"),
+                                               required=False)
+    resulting_check_date = forms.DateField(
         initial=get_now, label=_(u"Check date"), widget=DatePicker)
 
     TYPES = [
-        FieldType('material_type', models.MaterialType, is_multiple=True),
-        FieldType('material_type_quality', models.MaterialTypeQualityType),
-        FieldType('object_type', models.ObjectType, is_multiple=True),
-        FieldType('object_type_quality', models.ObjectTypeQualityType),
-        FieldType('communicabilitie', models.CommunicabilityType,
+        FieldType('resulting_material_type', models.MaterialType,
+                  is_multiple=True),
+        FieldType('resulting_material_type_quality',
+                  models.MaterialTypeQualityType),
+        FieldType('resulting_object_type', models.ObjectType,
+                  is_multiple=True),
+        FieldType('resulting_object_type_quality',
+                  models.ObjectTypeQualityType),
+        FieldType('resulting_communicabilitie', models.CommunicabilityType,
+                  is_multiple=True),
+        FieldType('resulting_checked_type', models.CheckedType,
                   is_multiple=True),
-        FieldType('checked_type', models.CheckedType, is_multiple=True),
     ]
 
 
diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 160e9d32d..4e2eb0faf 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -71,7 +71,7 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         "person__cached_label": _(u"Responsible"),
     }
     # extra keys than can be passed to save method
-    EXTRA_SAVED_KEYS = ('items', 'user')
+    EXTRA_SAVED_KEYS = ('items', 'user', 'resulting_find', 'upstream_items')
 
     # alternative names of fields for searches
     ALT_NAMES = {
@@ -292,10 +292,73 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         if q.count():
             self.index = q.all().aggregate(Max('index'))['index__max'] + 1
 
+    def _create_resulting_find(self, resulting_find, upstream_items):
+        m2m = {}
+        base_fields = [f.name for f in Find._meta.get_fields()]
+        for k in resulting_find.keys():
+            # if not in base fields should be a m2m
+            if k not in base_fields:
+                values = resulting_find.pop(k)
+                if values:
+                    m2m[k + "s"] = values
+
+        resulting_find['history_modifier'] = self.history_modifier
+        new_find = Find.objects.create(**resulting_find)
+
+        for k in m2m:
+            m2m_field = getattr(new_find, k)
+            try:
+                for value in m2m[k]:
+                    m2m_field.add(value)
+            except TypeError:
+                m2m_field.add(m2m[k])
+
+        dating_keys = ["period", "start_date", "end_date", "dating_type",
+                       "quality", "precise_dating"]
+        current_datings = []
+        current_base_finds = []
+        for upstream_item in upstream_items:
+            # datings are not explicitly part of the resulting_find
+            # need to reassociate with no duplicate
+            for dating in upstream_item.datings.all():
+                current_dating = []
+                for key in dating_keys:
+                    value = getattr(dating, key)
+                    if hasattr(value, 'pk'):
+                        value = value.pk
+                    current_dating.append(value)
+                current_dating = tuple(current_dating)
+                if current_dating in current_datings:
+                    # do not add similar dating
+                    continue
+                current_datings.append(current_dating)
+                dating.pk = None  # duplicate
+                dating.save()
+                new_find.datings.add(dating)
+
+            # associate base finds
+            for base_find in upstream_item.base_finds.all():
+                if base_find.pk in current_base_finds:
+                    continue
+                current_base_finds.append(base_find.pk)
+                new_find.base_finds.add(base_find)
+
+            upstream_item.downstream_treatment = self
+            upstream_item.history_modifier = self.history_modifier
+            upstream_item.save()
+        new_find.upstream_treatment = self
+        new_find.skip_history_when_saving = True
+        new_find.save()
+
     def save(self, *args, **kwargs):
-        items, user, extra_args_for_new = [], None, []
+        items, user, extra_args_for_new, resulting_find = [], None, [], None
+        upstream_items = []
         if "items" in kwargs:
             items = kwargs.pop('items')
+        if "resulting_find" in kwargs:
+            resulting_find = kwargs.pop('resulting_find')
+        if "upstream_items" in kwargs:
+            upstream_items = kwargs.pop('upstream_items')
         if "user" in kwargs:
             user = kwargs.pop('user')
         if "extra_args_for_new" in kwargs:
@@ -303,8 +366,15 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         self.pre_save()
         super(Treatment, self).save(*args, **kwargs)
         updated = []
+        # baskets
         if hasattr(items, "items"):
             items = items.items.all()
+        if hasattr(upstream_items, "items"):
+            upstream_items = upstream_items.items.all()
+
+        if upstream_items and resulting_find:
+            self._create_resulting_find(resulting_find, upstream_items)
+            return
 
         tps = list(self.treatment_types.all())
         has_destructive_tp = bool([tp for tp in tps if tp.destructive])
diff --git a/archaeological_finds/views.py b/archaeological_finds/views.py
index 789c420de..d6924257f 100644
--- a/archaeological_finds/views.py
+++ b/archaeological_finds/views.py
@@ -449,7 +449,7 @@ treatment_n1_wizard_steps = [
     ('final-treatment_creation_n1', FinalForm)
 ]
 
-treatment_creation_n1_wizard = TreatmentWizard.as_view(
+treatment_creation_n1_wizard = TreatmentN1Wizard.as_view(
     treatment_n1_wizard_steps,
     label=_(u"New treatment"),
     url_name='treatment_creation_n1',)
diff --git a/archaeological_finds/wizards.py b/archaeological_finds/wizards.py
index 8a7a3a513..3043d22aa 100644
--- a/archaeological_finds/wizards.py
+++ b/archaeological_finds/wizards.py
@@ -106,24 +106,16 @@ class TreatmentSearch(SearchWizard):
     model = models.Treatment
 
 
-class TreatmentWizard(Wizard):
+class TreatmentBase(Wizard):
     model = models.Treatment
     wizard_done_window = reverse_lazy('show-treatment')
-    basket_step = 'basetreatment-treatment_creation'
-    saved_args = {"items": []}
-
-    def get_form_kwargs(self, step, **kwargs):
-        kwargs = super(TreatmentWizard, self).get_form_kwargs(step, **kwargs)
-        if self.basket_step not in step:
-            return kwargs
-        kwargs['user'] = self.request.user
-        return kwargs
+    base_url = ""
 
     def get_current_finds(self):
         step = self.steps.current
         if not step:
             return
-        find_form_key = 'selecfind-treatment_creation'
+        find_form_key = 'selecfind-' + self.base_url
         find_ids = self.session_get_value(find_form_key, "resulting_pk")
         try:
             return [
@@ -134,8 +126,9 @@ class TreatmentWizard(Wizard):
             pass
 
     def get_form_initial(self, step, data=None):
-        initial = super(TreatmentWizard, self).get_form_initial(step)
-        if step != 'basetreatment-treatment_creation':
+        initial = super(TreatmentBase, self).get_form_initial(step)
+        base_step = 'basetreatment-' + self.base_url
+        if step != base_step:
             return initial
         finds = self.get_current_finds()
         if not finds:
@@ -150,6 +143,19 @@ class TreatmentWizard(Wizard):
         initial['location'] = locations[0]
         return initial
 
+
+class TreatmentWizard(TreatmentBase):
+    basket_step = 'basetreatment-treatment_creation'
+    saved_args = {"items": []}
+    base_url = 'treatment_creation'
+
+    def get_form_kwargs(self, step, **kwargs):
+        kwargs = super(TreatmentWizard, self).get_form_kwargs(step, **kwargs)
+        if self.basket_step not in step:
+            return kwargs
+        kwargs['user'] = self.request.user
+        return kwargs
+
     def get_extra_model(self, dct, form_list):
         """
         Get items concerned by the treatment
@@ -181,6 +187,121 @@ class TreatmentModificationWizard(TreatmentWizard):
     modification = True
 
 
+class TreatmentN1Wizard(TreatmentBase):
+    saved_args = {"upstream_items": [], "resulting_find": None}
+    base_url = 'treatment_creation_n1'
+
+    def _update_simple_initial_from_finds(self, initial, find, k):
+        r_k = "resulting_" + k
+        if getattr(find, k) and r_k in initial:
+            # pop k when value is inconsistent between finds
+            if initial[r_k] and initial[r_k] != getattr(find, k).pk:
+                initial.pop(r_k)
+            else:
+                initial[r_k] = getattr(find, k).pk
+        return initial
+
+    def _update_multi_initial_from_finds(self, initial, find, k):
+        r_k = "resulting_" + k
+        for value in getattr(find, k + 's').all():
+            if value.pk not in initial[r_k]:
+                initial[r_k].append(value.pk)
+        return initial
+
+    def _update_num_initial_from_finds(self, initial, find, k):
+        r_k = "resulting_" + k
+        if not getattr(find, k):
+            return initial
+        if initial[r_k] is None:
+            initial[r_k] = 0
+        initial[r_k] += getattr(find, k)
+        return initial
+
+    def _update_char_initial_from_finds(self, initial, find, k, sep=' ; '):
+        r_k = "resulting_" + k
+        if not getattr(find, k):
+            return initial
+        if not initial[r_k]:
+            initial[r_k] = getattr(find, k)
+        else:
+            initial[r_k] += sep + getattr(find, k)
+        return initial
+
+    def get_form_initial(self, step, data=None):
+        initial = super(TreatmentN1Wizard, self).get_form_initial(step)
+        if step != 'resultingfind-treatment_creation_n1':
+            return initial
+        finds = self.get_current_finds()
+        if not finds:
+            return initial
+        simple_key = ['material_type_quality']
+        multi_key = ['material_type', 'object_type', 'communicabilitie']
+        numeric_key = ['find_number', 'min_number_of_individuals']
+        desc_key = ['decoration', 'inscription', 'comment', 'dating_comment']
+        char_key = ['manufacturing_place']
+
+        for k in simple_key + numeric_key + desc_key + char_key:
+            initial["resulting_" + k] = None
+        for k in multi_key:
+            initial["resulting_" + k] = []
+
+        for find in finds:
+            for k in simple_key:
+                initial = self._update_simple_initial_from_finds(
+                    initial, find, k)
+            for k in multi_key:
+                initial = self._update_multi_initial_from_finds(
+                    initial, find, k)
+            for k in numeric_key:
+                initial = self._update_num_initial_from_finds(
+                    initial, find, k)
+            for k in char_key:
+                initial = self._update_char_initial_from_finds(
+                    initial, find, k, sep=' ; ')
+            for k in desc_key:
+                initial = self._update_char_initial_from_finds(
+                    initial, find, k, sep='\n')
+
+        for k in initial.keys():
+            if initial[k] is None:
+                initial.pop(k)
+        return initial
+
+    def get_extra_model(self, dct, form_list):
+        """
+        Get items concerned by the treatment
+        """
+        dct = super(TreatmentN1Wizard, self).get_extra_model(dct, form_list)
+        if 'resulting_pk' not in dct:
+            return dct
+
+        dct['upstream_items'] = []
+        # manage upstream items
+        pks = dct.pop('resulting_pk')
+        if hasattr(pks, 'split'):
+            pks = pks.split(u',')  # unicode or string
+        for pk in pks:
+            try:
+                find = models.Find.objects.get(pk=pk)
+                dct['upstream_items'].append(find)
+            except models.Find.DoesNotExist:
+                raise PermissionDenied
+
+        for find in dct['upstream_items']:
+            if 'own' in self.current_right \
+                    and not find.is_own(dct['history_modifier']):
+                raise PermissionDenied
+
+        # extract data of the new find
+        dct['resulting_find'] = {}
+        for k in dct.keys():
+            if k.startswith('resulting_') and k != "resulting_find":
+                dct['resulting_find'][
+                    k[len('resulting_'):]
+                ] = dct.pop(k)
+        return dct
+
+
 class TreatmentDeletionWizard(DeletionWizard):
     model = models.Treatment
     fields = ['label', 'other_reference', 'year', 'index',