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

def conv(i, k, s, a = None, t = (1, 1)):
    l1 = kl.Conv2D(k, s, strides = t, padding = 'same')(i)
    l2 = kl.BatchNormalization()(l1)

    if a:
        return kl.Activation(a)(l2)
    else:
        return l2

def sepc(i, k, s, a = None, t = (1, 1)):
    l1 = kl.SeparableConv2D(k, s, strides = t, padding = 'same')(i)
    l2 = kl.BatchNormalization()(l1)

    if a:
        return kl.Activation(a)(l2)
    else:
        return l2

def maxp(i, s):
    return kl.MaxPooling2D(s, padding = 'same')(i)

def addl(*i):
    return kl.Add()([*i])

def gapl(i):
    return kl.GlobalAveragePooling2D()(i)

def drpt(i, r):
    return kl.Dropout(r)(i)

def dnsl(i, s):
    return kl.Dense(s)(i)

def rscl(i, r):
    return kl.Rescaling(r)(i)


def raise_xception(si = (300, 300), cc = 1):
    i = kl.Input(shape = (*si, 3))

    b01 = rscl(i, 1/255)
    
    b11 = conv(b01, 16, (3,3), 'relu')

    b21 = conv(b11, 32, (1,1), t = (2,2))
    b22 = sepc(b11, 32, (3,3), 'relu')
    b23 = maxp(b22, (2,2))

    b31 = addl(b21, b23)

    b41 = conv(b31, 64, (1,1), t = (2,2))
    b42 = sepc(b31, 64, (3,3), 'relu')
    b43 = maxp(b42, (2,2))

    b51 = addl(b41, b43)

    b61 = sepc(b51, 128, (3,3))
    b71 = gapl(b61)

    b81 = drpt(b71, 0.3)
    b91 = dnsl(b81, 64)

    o = dnsl(b91, cc)

    return km.Model(inputs = i, outputs = o)


m = raise_xception()

m.compile(
        optimizer = ko.Lion(learning_rate = 0.002),
        loss = ks.BinaryCrossentropy(from_logits = True),
        metrics = ['accuracy'])

m.summary()