Adapt fields/forms/wizard for treatments

1 files changed, 67 insertions, 28 deletions

diff --git a/archaeological_finds/models.py b/archaeological_finds/models.py
index 2957c656d..01c904cc8 100644
--- a/archaeological_finds/models.py
+++ b/archaeological_finds/models.py
@@ -30,7 +30,7 @@ from ishtar_common.utils import cached_label_changed
 
 from ishtar_common.models import GeneralType, ImageModel, BaseHistorizedItem, \
     ShortMenuItem, LightHistorizedItem, HistoricalRecords, OwnPerms, Source, \
-    Person, Basket, get_external_id, post_save_cache
+    Organization, Person, Basket, get_external_id, post_save_cache
 
 from archaeological_operations.models import AdministrativeAct
 from archaeological_context_records.models import ContextRecord, Dating
@@ -792,16 +792,16 @@ post_save.connect(post_save_cache, sender=TreatmentType)
 post_delete.connect(post_save_cache, sender=TreatmentType)
 
 
-class Treatment(BaseHistorizedItem, OwnPerms):
+class Treatment(BaseHistorizedItem, ImageModel, OwnPerms):
     SHOW_URL = 'show-treatment'
-    external_id = models.CharField(_(u"External ID"), blank=True, null=True,
-                                   max_length=120)
-    container = models.ForeignKey(Container, verbose_name=_(u"Container"),
-                                  blank=True, null=True)
-    description = models.TextField(_(u"Description"), blank=True, null=True)
-    comment = models.TextField(_(u"Comment"), blank=True, null=True)
-    treatment_type = models.ForeignKey(TreatmentType,
-                                       verbose_name=_(u"Treatment type"))
+    label = models.CharField(_(u"Label"), blank=True, null=True,
+                             max_length=200)
+    other_reference = models.CharField(_(u"Other ref."), blank=True, null=True,
+                                       max_length=200)
+    year = models.IntegerField(_(u"Year"), default=2016)
+    index = models.IntegerField(_(u"Index"), default=1)
+    treatment_types = models.ManyToManyField(
+        TreatmentType, verbose_name=_(u"Treatment type"))
     location = models.ForeignKey(
         Warehouse, verbose_name=_(u"Location"), blank=True, null=True,
         help_text=_(
@@ -810,12 +810,23 @@ class Treatment(BaseHistorizedItem, OwnPerms):
     other_location = models.CharField(_(u"Location (not referenced)"),
                                       max_length=200, blank=True, null=True)
     person = models.ForeignKey(
-        Person, verbose_name=_(u"Doer"), blank=True, null=True,
+        Person, verbose_name=_(u"Responsible"), blank=True, null=True,
+        on_delete=models.SET_NULL, related_name='treatments')
+    organization = models.ForeignKey(
+        Organization, verbose_name=_(u"Organization"), blank=True, null=True,
         on_delete=models.SET_NULL, related_name='treatments')
+    external_id = models.CharField(_(u"External ID"), blank=True, null=True,
+                                   max_length=200)
+    comment = models.TextField(_(u"Comment"), blank=True, null=True)
+    description = models.TextField(_(u"Description"), blank=True, null=True)
+    goal = models.TextField(_(u"Goal"), blank=True, null=True)
     start_date = models.DateField(_(u"Start date"), blank=True, null=True)
     end_date = models.DateField(_(u"End date"), blank=True, null=True)
+    container = models.ForeignKey(Container, verbose_name=_(u"Container"),
+                                  blank=True, null=True)
+    target_is_basket = models.BooleanField(_("Target a basket"), default=False)
     history = HistoricalRecords()
-    TABLE_COLS = ('treatment_type__label', 'person', 'start_date',
+    TABLE_COLS = ('treatment_type', 'person', 'start_date',
                   'downstream_cached_label', 'upstream_cached_label')
     EXTRA_REQUEST_KEYS = {
         "downstream_cached_label": "downstream__cached_label",
@@ -829,6 +840,7 @@ class Treatment(BaseHistorizedItem, OwnPerms):
     class Meta:
         verbose_name = _(u"Treatment")
         verbose_name_plural = _(u"Treatments")
+        unique_together = ('year', 'index')
         permissions = (
             ("view_treatment", ugettext(u"Can view all Treatments")),
             ("view_own_treatment", ugettext(u"Can view own Treatment")),
@@ -843,6 +855,22 @@ class Treatment(BaseHistorizedItem, OwnPerms):
             lbl += u" %s %s" % (_(u"by"), unicode(self.person))
         return lbl
 
+    def treatment_type(self):
+        """
+        Treatment types label
+        :return: string
+        """
+        return u" ; ".join([unicode(t) for t in self.treatment_types.all()])
+
+    def pre_save(self):
+        # is not new
+        if self.pk is not None:
+            return
+        self.index = 1
+        q = Treatment.objects.filter(year=self.year)
+        if q.count():
+            self.index = q.all().aggregate(Max('index'))['index__max'] + 1
+
     def save(self, *args, **kwargs):
         items, user, extra_args_for_new = [], None, []
         if "items" in kwargs:
@@ -852,23 +880,34 @@ class Treatment(BaseHistorizedItem, OwnPerms):
         if "extra_args_for_new" in kwargs:
             extra_args_for_new = kwargs.pop('extra_args_for_new')
         is_new = self.pk is None
+        self.pre_save()
         super(Treatment, self).save(*args, **kwargs)
-        if not is_new or not items:
-            return
-        if hasattr(items, "items"):
-            items = items.items.all()
-        for item in items:
-            new = item.duplicate(user)
-            item.downstream_treatment = self
-            item.save()
-            new.upstream_treatment = self
-            for k in extra_args_for_new:
-                setattr(new, k, extra_args_for_new[k])
-            new.save()
-            # update baskets
-            for basket in FindBasket.objects.filter(items__pk=item.pk).all():
-                basket.items.remove(item)
-                basket.items.add(new)
+        updated = []
+        if is_new:
+            if hasattr(items, "items"):
+                items = items.items.all()
+            for item in items:
+                new = item.duplicate(user)
+                item.downstream_treatment = self
+                item.save()
+                new.upstream_treatment = self
+                for k in extra_args_for_new:
+                    setattr(new, k, extra_args_for_new[k])
+                new.save()
+                updated.append(new.pk)
+                # update baskets
+                for basket in \
+                        FindBasket.objects.filter(items__pk=item.pk).all():
+                    basket.items.remove(item)
+                    basket.items.add(new)
+        # manage containers
+        for find in Find.objects.filter(upstream_treatment=self).all():
+            if find.container != self.container:
+                find.container = self.container
+                if find.pk in updated:
+                    # don't record twice history
+                    find.skip_history_when_saving = True
+                find.save()
 
 
 class AbsFindTreatments(models.Model):