import tensorflow as tf
import numpy as np

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

def generate_inception_model(a = 5, b = 4, c = 2):
    # start
    i = kl.Input(shape = (300, 300, 3))

    r1 = kl.Conv2D(1024, (3, 3), padding = 'same', activation = 'relu')(i)
    r2 = kl.Conv2D(1024, (5, 5), padding = 'same', activation = 'relu')(r1)
    r3 = kl.Conv2D(1024, (7, 7), padding = 'same', activation = 'relu')(r2)

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

    r4 = kl.Conv2D(1024, (3, 3), padding = 'same', activation = 'relu')(m1)
    r5 = kl.Conv2D(1024, (5, 5), padding = 'same', activation = 'relu')(r4)

    ia0 = kl.MaxPooling2D((2, 2))(r5)

    # a types
    for k in range(a):
        exec(f"ia{k}_1_1 = kl.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(ia{k})")

        exec(f"ia{k}_3_1 = kl.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(ia{k})")
        exec(f"ia{k}_3_2 = kl.Conv2D(1024, (3, 1), padding = 'same', activation = 'relu')(ia{k}_3_1)")

        exec(f"ia{k}_5_1 = kl.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(ia{k})")
        exec(f"ia{k}_5_2 = kl.Conv2D(1024, (3, 1), padding = 'same', activation = 'relu')(ia{k}_5_1)")
        exec(f"ia{k}_5_3 = kl.Conv2D(1024, (5, 1), padding = 'same', activation = 'relu')(ia{k}_5_2)")

        exec(f"ia{k}_7_1 = kl.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(ia{k})")
        exec(f"ia{k}_7_2 = kl.Conv2D(1024, (3, 1), padding = 'same', activation = 'relu')(ia{k}_7_1)")
        exec(f"ia{k}_7_3 = kl.Conv2D(1024, (5, 1), padding = 'same', activation = 'relu')(ia{k}_7_2)")
        exec(f"ia{k}_7_4 = kl.Conv2D(1024, (7, 1), padding = 'same', activation = 'relu')(ia{k}_7_3)")

        exec(f"ia{k+1} = kl.Concatenate()([ia{k}_1_1, ia{k}_3_2, ia{k}_5_3, ia{k}_7_4])")

    # grid size reductor 1
    iab_1 = kl.Conv2D(1024, (5, 1), padding = 'same', activation = 'relu')(eval(f"ia{a}"))

    iab_2 = kl.Conv2D(1024, (5, 1), padding = 'same', activation = 'relu')(eval(f"ia{a}"))
    iab_3 = kl.Conv2D(1024, (3, 1), padding = 'same', activation = 'relu')(iab_2)
    iab_4 = kl.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(iab_3)

    iab_6 = kl.Concatenate()([iab_1, iab_4, iab_5])

    # b types
    for k in range(b):
        exec(f"ib{k}_1_1 = kl.MaxPooling2D((2, 2), padding = 'same')(iab_6)")
        exec(f"ib{k}_1_2 = kl.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(i)")

        exec(f"ib{k}_3_1 = kl.Conv2D(1024, (3, 1), padding = 'same', activation = 'relu')(iab_6)")

        exec(f"ib{k}_5_1 = kl.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(iab_6)")
        exec(f"ib{k}_5_2 = kl.Conv2D(1024, (3, 1), padding = 'same', activation = 'relu')(ib{k}_5_1)")
        exec(f"ib{k}_5_3 = kl.Conv2D(1024, (5, 1), padding = 'same', activation = 'relu')(ib{k}_5_2)")

        exec(f"ib{k}_7_1 = kl.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(iab_6)")
        exec(f"ib{k}_7_2 = kl.Conv2D(1024, (3, 1), padding = 'same', activation = 'relu')(ib{k}_7_1)")
        exec(f"ib{k}_7_3 = kl.Conv2D(1024, (5, 1), padding = 'same', activation = 'relu')(ib{k}_7_2)")
        exec(f"ib{k}_7_4 = kl.Conv2D(1024, (7, 1), padding = 'same', activation = 'relu')(ib{k}_7_3)")

        exec(f"ib{k+1} = kl.Concatenate()([ib{k}_1_2, ib{k}_3_2, ib{k}_5_3, ib{k}_7_4])")

    # grid size reductor 2

    # c types
    for k in range(c):
        exec(f"ic{k}_1_1 = kl.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(i{'b' if k else 'a'}{k if k else a})")

        exec(f"ic{k}_3_1 = kl.Conv2D(1024, (3, 1), padding = 'same', activation = 'relu')(i{'b' if k else 'a'}{k if k else a})")

        exec(f"ic{k}_5_1 = kl.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(i{'b' if k else 'a'}{k if k else a})")
        exec(f"ic{k}_5_2 = kl.Conv2D(1024, (3, 1), padding = 'same', activation = 'relu')(ib{k}_5_1)")
        exec(f"ic{k}_5_3 = kl.Conv2D(1024, (5, 1), padding = 'same', activation = 'relu')(ib{k}_5_2)")

        exec(f"ic{k}_7_1 = kl.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(i{'b' if k else 'a'}{k if k else a})")
        exec(f"ic{k}_7_2 = kl.Conv2D(1024, (3, 1), padding = 'same', activation = 'relu')(ib{k}_7_1)")
        exec(f"ic{k}_7_3 = kl.Conv2D(1024, (5, 1), padding = 'same', activation = 'relu')(ib{k}_7_2)")
        exec(f"ic{k}_7_4 = kl.Conv2D(1024, (7, 1), padding = 'same', activation = 'relu')(ib{k}_7_3)")

        exec(f"ib{k+1} = kl.Concatenate()([ib{k}_1_1, ib{k}_3_2, ib{k}_5_3, ib{k}_7_4])")

    # mppt

    o = eval(f"ic{c}")
    return tf.keras.Model(inputs = i, outputs = o)


'''
a_i = l.Input(shape = (300, 300, 3))

a_1_1 = l.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(a_i)

a_3_1 = l.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(a_i)
a_3_2 = l.Conv2D(1024, (3, 1), padding = 'same', activation = 'relu')(a_3_1)

a_5_1 = l.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(a_i)
a_5_2 = l.Conv2D(1024, (5, 1), padding = 'same', activation = 'relu')(a_5_1)

a_7_1 = l.Conv2D(1024, (1, 1), padding = 'same', activation = 'relu')(a_i)
a_7_2 = l.Conv2D(1024, (7, 1), padding = 'same', activation = 'relu')(a_7_1)

a_o = l.Concatenate()([a_1_1, a_3_2, a_5_2, a_7_2])

inception_type_a = [a_i, a_o]

tf.keras.Model(*inception_type_a)
'''

mod = generate_inception_model()
mod.summary()
mod.compile(optimizer = ko.Lion(learning_rate = 0.001))