""" Contains metrics for segmentation """
import numpy as np

from . import ClassificationMetrics, get_components


class SegmentationMetricsByPixels(ClassificationMetrics):
    """ Metrics to assess segmentation models pixel-wise """
    pass


class SegmentationMetricsByInstances(ClassificationMetrics):
    """ Metrics to assess segmentation models by instances (connected components) """
    def __init__(self, targets, predictions, fmt='proba', num_classes=None, axis=None, threshold=.5, iot=.5):
        super().__init__(targets, predictions, fmt, num_classes, axis, threshold, confusion=False)
        self.iot = iot
        self.target_instances = self._get_instances(self.targets)
        self.predicted_instances = self._get_instances(self.predictions)
        self._calc_confusion()

    def _get_instances(self, inputs):
        return get_components(inputs, batch=True)

    def _calc_confusion(self):
        self._confusion_matrix = np.zeros((self.targets.shape[0], self.num_classes - 1, 2, 2), dtype=np.intp)

        for i in range(len(self.target_instances)):
            for j in range(len(self.target_instances[i])):
                for k in range(self.num_classes - 1):
                    c = self.target_instances[i][j]
                    coords = [slice(i, i+1)]
                    for d in range(c.shape[0]):
                        coords.append(slice(c[d, 0], c[d, 1]))
                    targ = self.targets[coords]
                    pred = self.targets[coords] * self.predictions[coords]
                    if np.sum(pred) / targ.size >= self.iot:
                        self._confusion_matrix[i, k, 1, 1] += 1
                    else:
                        self._confusion_matrix[i, k, 0, 1] += 1

        for i in range(len(self.predicted_instances)):
            for j in range(len(self.predicted_instances[i])):
                for k in range(self.num_classes - 1):
                    c = self.predicted_instances[i][j]
                    coords = [slice(i, i+1)]
                    for d in range(c.shape[0]):
                        coords.append(slice(c[d, 0], c[d, 1]))
                    pred = self.predictions[coords]
                    targ = self.targets[coords]
                    if np.sum(pred) / targ.size < self.iot:
                        self._confusion_matrix[i, k, 1, 0] += 1

    def true_positive(self, label=None, *args, **kwargs):
        _ = args, kwargs
        return self._count(lambda l: self._confusion_matrix[:, l-1, 1, 1], label)

    def true_negative(self, label=None, *args, **kwargs):
        raise ValueError("True negaitve is inapplicable for instance-based metrics")

    def condition_positive(self, label=None, *args, **kwargs):
        _ = args, kwargs
        return self._count(lambda l: self._confusion_matrix[:, l-1, :, 1].sum(axis=1), label)

    def prediction_positive(self, label=None, *args, **kwargs):
        _ = args, kwargs
        return self._count(lambda l: self._confusion_matrix[:, l-1, 1].sum(axis=1), label)

    def total_population(self, label=None, *args, **kwargs):
        _ = args, kwargs
        return self._return(self._confusion_matrix[:, label - 1].sum(axis=(1, 2)))

@@ -682,16 +682,22 @@
 .targets
+, axis
 )%0A      
@@ -761,16 +761,22 @@
 dictions
+, axis
 )%0A      
@@ -836,16 +836,22 @@
 , inputs
+, axis
 ):%0A     
@@ -853,22 +853,293 @@
         
-return
+if axis is None:%0A            instances = %5Bget_components(inputs, batch=True)%5D%0A        else:%0A            instances = %5B%5D%0A            shape = %5Bslice(None)%5D * inputs.ndim%0A            for i in range(1, inputs.shape%5Baxis%5D):%0A                shape%5Baxis%5D = i%0A                per_class =
  get_com
@@ -1144,32 +1144,39 @@
 omponents(inputs
+%5Bshape%5D
 , batch=True)%0A%0A 
@@ -1172,16 +1172,85 @@
 ch=True)
+%0A                instances.append(per_class)%0A        return instances
 %0A%0A    de
@@ -1389,33 +1389,33 @@
 p)%0A%0A        for 
-i
+k
  in range(len(se
@@ -1412,34 +1412,28 @@
 nge(
-len(
 self.
-target_instances)
+num_classes - 1
 ):%0A 
@@ -1439,33 +1439,33 @@
             for 
-j
+i
  in range(len(se
@@ -1476,33 +1476,33 @@
 arget_instances%5B
-i
+k
 %5D)):%0A           
@@ -1502,33 +1502,33 @@
             for 
-k
+j
  in range(self.n
@@ -1521,36 +1521,48 @@
 n range(
+len(
 self.
-num_classes - 1
+target_instances%5Bk%5D%5Bi%5D)
 ):%0A     
@@ -1594,32 +1594,35 @@
 arget_instances%5B
+k%5D%5B
 i%5D%5Bj%5D%0A          
@@ -2110,33 +2110,33 @@
  1%0A%0A        for 
-i
+k
  in range(len(se
@@ -2133,37 +2133,28 @@
 nge(
-len(
 self.
-predicted_instances)
+num_classes - 1
 ):%0A 
@@ -2160,33 +2160,33 @@
             for 
-j
+i
  in range(len(se
@@ -2208,17 +2208,17 @@
 stances%5B
-i
+k
 %5D)):%0A   
@@ -2226,33 +2226,33 @@
             for 
-k
+j
  in range(self.n
@@ -2245,36 +2245,51 @@
 n range(
+len(
 self.
-num_classes - 1
+predicted_instances%5Bk%5D%5Bi%5D)
 ):%0A     
@@ -2332,16 +2332,19 @@
 stances%5B
+k%5D%5B
 i%5D%5Bj%5D%0A  

Add multiclass segmentation