import tensorflow as tf

from autokeras.engine import block as block_module
from autokeras.engine import io_hypermodel


class Head(block_module.Block, io_hypermodel.IOHyperModel):
    """Base class for the heads, e.g. classification, regression.

    # Arguments
        loss: A Keras loss function. Defaults to None. If None, the loss will be
            inferred from the AutoModel.
        metrics: A list of Keras metrics. Defaults to None. If None, the metrics will
            be inferred from the AutoModel.
        output_shape: Tuple of int(s). Defaults to None. If None, the output shape
            will be inferred from the AutoModel.
    """

    def __init__(self, loss=None, metrics=None, output_shape=None, **kwargs):
        super().__init__(**kwargs)
        self.output_shape = output_shape
        self.loss = tf.keras.losses.get(loss)
        if metrics is None:
            metrics = []
        self.metrics = [tf.keras.metrics.get(metric) for metric in metrics]
        # Mark if the head should directly output the input tensor.

    def get_config(self):
        config = super().get_config()
        config.update({
            'loss': tf.keras.losses.serialize(self.loss),
            'metrics': [tf.keras.metrics.serialize(metric)
                        for metric in self.metrics],
            'output_shape': self.output_shape
        })
        return config

    @classmethod
    def from_config(cls, config):
        config['loss'] = tf.keras.losses.deserialize(config['loss'])
        config['metrics'] = [tf.keras.metrics.deserialize(metric)
                             for metric in config['metrics']]
        return super().from_config(config)

    def build(self, hp, inputs=None):
        raise NotImplementedError

    def config_from_adapter(self, adapter):
        self.output_shape = adapter.shape

@@ -114,16 +114,800 @@
 model%0A%0A%0A
+def serialize_metrics(metrics):%0A    serialized = %5B%5D%0A    for metric in metrics:%0A        if isinstance(metric, str):%0A            serialized.append(%5Bmetric%5D)%0A        else:%0A            serialized.append(tf.keras.metrics.serialize(metric))%0A    return serialized%0A%0A%0Adef deserialize_metrics(metrics):%0A    deserialized = %5B%5D%0A    for metric in metrics:%0A        if isinstance(metric, list):%0A            deserialized.append(metric%5B0%5D)%0A        else:%0A            deserialized.append(tf.keras.metrics.deserialize(metric))%0A    return deserialized%0A%0A%0Adef serialize_loss(loss):%0A    if isinstance(loss, str):%0A        return %5Bloss%5D%0A    return tf.keras.losses.serialize(loss)%0A%0A%0Adef deserialize_loss(loss):%0A    if isinstance(loss, list):%0A        return loss%5B0%5D%0A    return tf.keras.losses.deserialize(loss)%0A%0A%0A
 class He
@@ -1612,33 +1612,12 @@
 s = 
-tf.keras.losses.get(loss)
+loss
 %0A   
@@ -1693,60 +1693,15 @@
 s = 
-%5Btf.keras.metrics.get(metric) for metric in metrics%5D
+metrics
 %0A   
@@ -1874,32 +1874,16 @@
 'loss': 
-tf.keras.losses.
 serializ
@@ -1883,16 +1883,21 @@
 erialize
+_loss
 (self.lo
@@ -1924,34 +1924,16 @@
 trics': 
-%5Btf.keras.metrics.
 serializ
@@ -1937,55 +1937,17 @@
 lize
-(
+_
 metric
-)%0A                        for metric in 
+s(
 self
@@ -1954,17 +1954,17 @@
 .metrics
-%5D
+)
 ,%0A      
@@ -2113,32 +2113,16 @@
 oss'%5D = 
-tf.keras.losses.
 deserial
@@ -2124,16 +2124,21 @@
 erialize
+_loss
 (config%5B
@@ -2178,26 +2178,8 @@
 %5D = 
-%5Btf.keras.metrics.
 dese
@@ -2189,60 +2189,17 @@
 lize
-(
+_
 metric
-)%0A                             for metric in 
+s(
 conf
@@ -2211,17 +2211,17 @@
 etrics'%5D
-%5D
+)
 %0A       

use user passed format for metrics and losses