import numpy as np, cv2, os
import matplotlib.pyplot as plt, seaborn as sns
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from sklearn.preprocessing import LabelEncoder
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense

def load_images(path, size):
    I,L=[],[]
    for c in os.listdir(path):
        cp=os.path.join(path,c)
        if not os.path.isdir(cp): continue
        for f in os.listdir(cp):
            if f.lower().endswith(('.jpg','.jpeg','.png')):
                img=cv2.imread(os.path.join(cp,f))
                if img is not None:
                    I.append(cv2.resize(img,size)/255.0)
                    L.append(c)
    return np.array(I,dtype=np.float32), np.array(L)

def get_arch(name,size,nc):
    s=(*size,3); r='relu'
    C,M,F,D=Conv2D,MaxPooling2D,Flatten,Dense
    return {
        "alexnet":[
            C(96,(11,11),strides=4,activation=r,input_shape=s),
            M((3,3),strides=2),
            C(256,(5,5),padding='same',activation=r),M((3,3),strides=2),
            C(384,(3,3),padding='same',activation=r),
            C(384,(3,3),padding='same',activation=r),
            C(256,(3,3),padding='same',activation=r),
            M((3,3),strides=2),F(),D(4096,r),D(4096,r),D(nc,'softmax')],
        "deep_cnn":[
            C(96,(7,7),strides=4,activation=r,input_shape=s),
            M((3,3),strides=2),
            C(256,(5,5),padding='same',activation=r),M((3,3),strides=2),
            C(512,(3,3),padding='same',activation=r),
            C(1024,(3,3),padding='same',activation=r),
            C(512,(3,3),padding='same',activation=r),
            M((3,3),strides=2),F(),D(4096,r),D(4096,r),D(nc,'softmax')],
        "light_cnn":[
            C(96,(7,7),strides=4,activation=r,input_shape=s),
            M((3,3),strides=2),
            C(256,(5,5),padding='same',activation=r),M((3,3),strides=2),
            F(),D(4096,r),D(4096,r),D(nc,'softmax')],
        "minimal_cnn":[
            C(96,(7,7),strides=4,activation=r,input_shape=s),
            M((3,3),strides=2),F(),D(50,r),D(nc,'softmax')]
    }[name]

def plot_history(h,title):
    plt.figure(figsize=(12,4))
    for i,k in enumerate(['accuracy','loss']):
        plt.subplot(1,2,i+1)
        plt.plot(h.history[k],label='Train')
        plt.plot(h.history[f'val_{k}'],label='Val')
        plt.title(f'{k} - {title}'); plt.legend()
    plt.tight_layout(); plt.show()

def evaluate(model,Xv,Yv,le,title):
    yp=np.argmax(model.predict(Xv),axis=1)
    acc=accuracy_score(Yv,yp)*100
    print(f"\n{'='*50}\n  {title}\n{'='*50}\nTest Accuracy: {acc:.2f}%")
    print("\nClassification Report:")
    print(classification_report(Yv,yp,target_names=le.classes_))
    plt.figure(figsize=(6,5))
    sns.heatmap(confusion_matrix(Yv,yp),annot=True,fmt='d',cmap='Blues',
                xticklabels=le.classes_,yticklabels=le.classes_)
    plt.title(f'CM - {title}'); plt.tight_layout(); plt.show()
    return acc

def predict_image(model,le,p,size):
    img=cv2.imread(p)
    if img is None: return print(f"Not found: {p}")
    rs=cv2.resize(img,size)
    yp=model.predict(np.expand_dims(rs/255.,0).astype(np.float32))
    c,cf=le.classes_[np.argmax(yp)],np.max(yp)*100
    print(f"\nImage: {p}\nPredicted: {c}\nConfidence: {cf:.2f}%")
    plt.figure(figsize=(4,4))
    plt.imshow(cv2.cvtColor(rs,cv2.COLOR_BGR2RGB))
    plt.title(f"{c} ({cf:.2f}%)"); plt.axis('off'); plt.show()

def run_experiment(train_path,test_path,models,epochs,image_size=(150,150),test_images=[]):
    Xt,yt=load_images(train_path,image_size)
    Xv,yv=load_images(test_path,image_size)
    print(f"Train:{Xt.shape} | Test:{Xv.shape}")
    le=LabelEncoder()
    Yt=le.fit_transform(yt).astype(np.int32)
    Yv=le.transform(yv).astype(np.int32)
    print(f"Classes: {le.classes_}")
    nc=len(le.classes_)
    labels={
        "alexnet"    :"Model 1: AlexNet",
        "deep_cnn"   :"Model 2: Deeper CNN",
        "light_cnn"  :"Model 3: Lighter CNN",
        "minimal_cnn":"Model 4: Minimal CNN"
    }
    R,TR={},{}
    for name,ep in zip(models,epochs):
        title=labels.get(name,name)
        print(f"\n>>> Training {title}")
        m=Sequential(get_arch(name,image_size,nc))
        m.compile('adam','sparse_categorical_crossentropy',['accuracy'])
        m.summary()
        h=m.fit(Xt,Yt,validation_data=(Xv,Yv),epochs=ep,verbose=1)
        plot_history(h,title)
        R[title]=evaluate(m,Xv,Yv,le,title)
        TR[title]=m
    best=max(R,key=R.get)
    print("\n"+"="*45)
    print("       FINAL MODEL COMPARISON SUMMARY")
    print("="*45)
    [print(f"  {n:<22} {a:.2f}%") for n,a in R.items()]
    print(f"\n  Best: {best} ({R[best]:.2f}%)")
    print("="*45)
    plt.figure(figsize=(10,5))
    plt.bar(R.keys(),R.values())
    [plt.text(i,a+1,f"{a:.2f}%",ha='center') for i,(_,a) in enumerate(R.items())]
    plt.title('Model Accuracy Comparison')
    plt.ylim(0,100); plt.tight_layout(); plt.show()
    print(f"\nPredicting using Best Model: {best}")
    for p in test_images:
        predict_image(TR[best],le,p,image_size) if os.path.exists(p) else print(f"Not found:{p}")
