import tensorflow as tf
import numpy as np

from tensorflow import keras as k
from tensorflow.keras import layers as kl
from tensorflow.keras import models as km
from tensorflow.keras import optimizers as ko
from tensorflow.keras import callbacks as kc

CONV_SIZE = 96
MPPT_SIZE = 256
DRPT_RATE = 0.3
TOUT_AMNT = 2

def conv(i, n, s, t = (1, 1)):
    c = kl.Conv2D(n, s, padding = 'same', strides = t)(i)
    b = kl.BatchNormalization()(c)
    return kl.Activation('relu')(b)

def avg(i, s):
    return kl.AveragePooling2D(s, padding = 'same', strides = (1, 1))(i)

def max(i, s):
    return kl.MaxPooling2D(s, padding = 'same', strides = (1, 1))(i)

def generate_start(i):
    s = kl.Rescaling(1./255.)(i)
    r1 = conv(s, CONV_SIZE, (3, 3))
    r2 = conv(r1, CONV_SIZE, (3, 3))
    r3 = conv(r2, CONV_SIZE, (3, 3))

    m1 = kl.MaxPooling2D((3, 3))(r3)

    r4 = conv(m1, CONV_SIZE, (1, 1))
    r5 = conv(r4, CONV_SIZE, (3, 3))

    return kl.MaxPooling2D((3, 3))(r5)

def generate_type_a(i):
    a11 = conv(i, CONV_SIZE, (1, 1))

    a21 = avg(i, (3, 3))
    a22 = conv(a21, CONV_SIZE, (1, 1))

    a31 = conv(i, CONV_SIZE, (1, 1))
    a32 = conv(a31, CONV_SIZE, (3, 3))

    a41 = conv(i, CONV_SIZE, (1, 1))
    a42 = conv(a41, CONV_SIZE, (3, 3))
    a43 = conv(a42, CONV_SIZE, (3, 3))

    return kl.Concatenate()([a11, a22, a32, a43])

def generate_ab_bridge(i):
    ab11 = conv(i, CONV_SIZE, (1, 1))
    ab12 = conv(ab11, CONV_SIZE, (3, 3))
    ab13 = conv(ab12, CONV_SIZE, (3, 3))

    ab21 = conv(i, CONV_SIZE, (3, 3))

    ab31 = max(i, (3, 3))

    return kl.Concatenate()([ab13, ab21, ab31])

def generate_type_b(i):
    b11 = conv(i, CONV_SIZE, (1, 1))
    b12 = conv(b11, CONV_SIZE, (7, 1))
    b13 = conv(b12, CONV_SIZE, (1, 7))
    b14 = conv(b13, CONV_SIZE, (7, 1))
    b15 = conv(b14, CONV_SIZE, (1, 7))

    b21 = conv(i, CONV_SIZE, (1, 1))
    b22 = conv(b21, CONV_SIZE, (1, 7))
    b23 = conv(b22, CONV_SIZE, (7, 1))

    b31 = conv(i, CONV_SIZE, (1, 1))

    b41 = avg(i, (3, 3))
    b42 = conv(b41, CONV_SIZE, (1, 1))

    return kl.Concatenate()([b15, b23, b31, b42])

def generate_bc_bridge(i):
    bc11 = conv(i, CONV_SIZE, (1, 1))
    bc12 = conv(bc11, CONV_SIZE, (7, 1))
    bc13 = conv(bc12, CONV_SIZE, (1, 7))
    bc14 = conv(bc13, CONV_SIZE, (3, 3))

    bc21 = conv(i, CONV_SIZE, (1, 1))
    bc22 = conv(bc21, CONV_SIZE, (1, 1))

    bc31 = max(i, (3, 3))

    return kl.Concatenate()([bc14, bc22, bc31])

def generate_aux(i):
    u1 = kl.AveragePooling2D((5, 5))(i)
    u2 = conv(u1, CONV_SIZE, (1, 1))
    u3 = kl.Flatten()(u2)
    u4 = kl.Dropout(DRPT_RATE)(u3)
    u5 = kl.Dense(MPPT_SIZE)(u4)
    u6 = kl.Dense(TOUT_AMNT)(u5)

    return kl.Activation('softmax')(u6)

def generate_type_c(i):
    c11 = conv(i, CONV_SIZE, (1, 1))
    c12 = conv(c11, CONV_SIZE, (3, 3))
    c13 = conv(c12, CONV_SIZE, (1, 3))
    c14 = conv(c12, CONV_SIZE, (3, 1))

    c21 = conv(i, CONV_SIZE, (1, 1))
    c22 = conv(c21, CONV_SIZE, (1, 3))
    c23 = conv(c21, CONV_SIZE, (3, 1))

    c31 = max(i, (3, 3))
    c32 = conv(c31, CONV_SIZE, (1, 1))

    c41 = conv(i, CONV_SIZE, (1, 1))

    return kl.Concatenate()([c14, c23, c32, c41])

def generate_finish(i):
    f1 = kl.AveragePooling2D((5, 5))(i)
    f2 = conv(f1, CONV_SIZE, (1, 1))
    f3 = kl.Flatten()(f2)
    f4 = kl.Dropout(DRPT_RATE)(f3)
    f5 = kl.Dense(MPPT_SIZE)(f4)
    f6 = kl.Dense(TOUT_AMNT)(f5)

    return kl.Activation('softmax')(f6)


gi = kl.Input((300, 300, 3))

ds = generate_start(gi)

for _ in range(3):
    ds = generate_type_a(ds)

ds = generate_ab_bridge(ds)

for _ in range(4):
    ds = generate_type_b(ds)

uo = generate_aux(ds)

go = generate_bc_bridge(ds)

for _ in range(2):
    go = generate_type_c(go)

go = generate_finish(go)

mod = k.Model(inputs = gi, outputs = go)

mod.summary()
mod.compile(optimizer = ko.Lion(learning_rate = 0.0001),
            loss = 'categorical_crossentropy',
            metrics = ['accuracy'])