from PIL import Image
import numpy as np
import os
import pickle
import time

# ==================== PATHS ====================
OUTPUT_FOLDER = "/storage/emulated/0/Download/"
MODEL_FOLDER = "/storage/emulated/0/AI_Photos/models/"

print("="*60)
print("🎨 IMAGE GENERATOR (FAST LAUNCH)")
print("="*60)

# ==================== NEURAL NETWORK CLASS ====================
class MyNeuralNetwork:
    def __init__(self, input_size, hidden_size, output_size):
        self.W1 = None
        self.b1 = None
        self.W2 = None
        self.b2 = None
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
    
    def forward(self, X):
        z1 = np.dot(X, self.W1) + self.b1
        a1 = np.tanh(z1)
        z2 = np.dot(a1, self.W2) + self.b2
        a2 = 1 / (1 + np.exp(-z2))
        return a2

# ==================== LOAD MODEL ====================
def load_model():
    """Loads model from folder"""
    model_paths = [
        os.path.join(MODEL_FOLDER, "model.pkl"),
        os.path.join(MODEL_FOLDER, "model.npz"),
        os.path.join(MODEL_FOLDER, "model.tflite")
    ]
    
    # Check for models
    found = False
    for path in model_paths:
        if os.path.exists(path):
            found = True
            print(f"📂 Found model: {os.path.basename(path)}")
    
    if not found:
        print("\n❌ MODEL NOT FOUND!")
        print(f"📁 Look in: {MODEL_FOLDER}")
        print("\n📌 First run the main script to train the model.")
        return None, None
    
    # Load pickle model (priority)
    pickle_path = os.path.join(MODEL_FOLDER, "model.pkl")
    if os.path.exists(pickle_path):
        try:
            print("🔄 Loading pickle model...")
            with open(pickle_path, 'rb') as f:
                data = pickle.load(f)
            
            nn = MyNeuralNetwork(data['input_size'], data['hidden_size'], data['output_size'])
            nn.W1 = data['W1']
            nn.b1 = data['b1']
            nn.W2 = data['W2']
            nn.b2 = data['b2']
            
            print("✅ Model loaded!")
            print(f"📐 Image size: {data.get('image_size', 40)}x{data.get('image_size', 40)}")
            return nn, data.get('image_size', 40)
        except Exception as e:
            print(f"❌ Error loading pickle: {e}")
    
    # Try loading numpy model
    numpy_path = os.path.join(MODEL_FOLDER, "model.npz")
    if os.path.exists(numpy_path):
        try:
            print("🔄 Loading numpy model...")
            data = np.load(numpy_path, allow_pickle=True)
            
            nn = MyNeuralNetwork(int(data['input_size']), int(data['hidden_size']), int(data['output_size']))
            nn.W1 = data['W1']
            nn.b1 = data['b1']
            nn.W2 = data['W2']
            nn.b2 = data['b2']
            
            print("✅ Model loaded!")
            return nn, 40
        except Exception as e:
            print(f"❌ Error loading numpy: {e}")
    
    # Try TFLite
    tflite_path = os.path.join(MODEL_FOLDER, "model.tflite")
    if os.path.exists(tflite_path):
        try:
            import tensorflow as tf
            print("🔄 Loading TFLite model...")
            interpreter = tf.lite.Interpreter(model_path=tflite_path)
            interpreter.allocate_tensors()
            print("✅ TFLite model loaded!")
            return interpreter, 40
        except Exception as e:
            print(f"❌ Error loading TFLite: {e}")
    
    return None, None

# ==================== GENERATION ====================
def generate_with_numpy(nn, size=40):
    """Generates image using NumPy model"""
    noise = np.random.randn(1, size*size*3) * 2.5
    output = nn.forward(noise)
    return output

def generate_with_tflite(interpreter, size=40):
    """Generates image using TFLite model"""
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    
    noise = np.random.randn(1, size*size*3).astype(np.float32) * 2.5
    interpreter.set_tensor(input_details[0]['index'], noise)
    interpreter.invoke()
    output = interpreter.get_tensor(output_details[0]['index'])
    return output

def save_image(array, filename, size=40):
    """Saves image"""
    arr = np.clip(array.reshape(size, size, 3), 0, 1) * 255
    img = Image.fromarray(arr.astype(np.uint8), mode='RGB')
    img.save(filename)
    return img

def show_preview(array, size=40):
    """Shows preview in console"""
    arr = np.clip(array.reshape(size, size, 3), 0, 1)
    symbols = ['  ', '░░', '▒▒', '▓▓', '██']
    print("\n👀 PREVIEW:")
    for i in range(0, size, 2):
        line = ''
        for j in range(0, size, 2):
            bright = int(np.mean(arr[i, j]) * 5)
            line += symbols[min(bright, 4)]
        print(line)

# ==================== MAIN FUNCTION ====================
def main():
    # 1. Load model
    model, size = load_model()
    if model is None:
        return
    
    # 2. Detect model type
    is_tflite = hasattr(model, 'get_input_details')
    
    print(f"\n🎨 GENERATING IMAGES...")
    print(f"📐 Size: {size}x{size}")
    print(f"🔄 Mode: {'TFLite (fast)' if is_tflite else 'NumPy'}")
    
    # 3. Generate images
    start_time = time.time()
    generated = []
    
    for i in range(4):
        if is_tflite:
            output = generate_with_tflite(model, size)
        else:
            output = generate_with_numpy(model, size)
        
        path = os.path.join(OUTPUT_FOLDER, f"generated_{i+1}.png")
        save_image(output, path, size)
        generated.append(output)
        print(f"  ✅ {path}")
    
    # 4. Show preview
    show_preview(generated[0], size)
    
    # 5. Create collage
    print("\n🎨 CREATING COLLAGE...")
    collage_size = size * 2
    collage = np.zeros((collage_size, collage_size, 3))
    positions = [(0,0), (0,size), (size,0), (size,size)]
    
    for i, (x, y) in enumerate(positions):
        img = np.clip(generated[i].reshape(size, size, 3), 0, 1) * 255
        collage[x:x+size, y:y+size] = img
    
    collage_img = Image.fromarray(collage.astype(np.uint8), mode='RGB')
    collage_path = os.path.join(OUTPUT_FOLDER, "collage_4_images.png")
    collage_img.save(collage_path)
    print(f"  ✅ {collage_path}")
    
    # 6. FINAL
    elapsed = time.time() - start_time
    print("\n" + "="*60)
    print(f"🎉 DONE! Generated in {elapsed:.1f} seconds")
    print(f"📁 Images in: {OUTPUT_FOLDER}")
    print("  - generated_1.png")
    print("  - generated_2.png")
    print("  - generated_3.png")
    print("  - generated_4.png")
    print("  - collage_4_images.png")
    print("\n📱 Open Gallery → Downloads folder")
    print("="*60)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n⏹️ Stopped by user")
    except Exception as e:
        print(f"\n❌ Error: {e}")
        print("\n📌 Make sure the model is trained and located in:")
        print(f"   {MODEL_FOLDER}")