import numpy as np
import cv2
import matplotlib.pyplot as plt
import cvzone
import gradio as gr
from sklearn.cluster import KMeans
from PIL import Image, ImageFont, ImageDraw, ImageColor
import os
import math


import boto3
import os

AWS_S3_CREDS = {
    "aws_access_key_id":os.environ['AccessKey'],
    "aws_secret_access_key":os.environ['SecretKey']
}
s3_client = boto3.client('s3',**AWS_S3_CREDS)

def get_color(image):
  pixels = image.reshape(-1, 3)

  num_colors = 2

  kmeans = KMeans(n_clusters=num_colors)
  kmeans.fit(pixels)

  dominant_colors = kmeans.cluster_centers_.astype(int)

  theme_color = dominant_colors[0]

  return theme_color

base_color = (195, 212, 220)

def generate_shades(base_color, num_shades):
    r, g, b = base_color

    shades = []

    for i in range(num_shades):
        brightness = float(i) / (num_shades - 1)
        new_r = int(r * brightness)
        new_g = int(g * brightness)
        new_b = int(b * brightness)
        shades.append((new_r, new_g, new_b))
    return shades

def text_to_image(
    text: str,
    font_filepath: str,
    font_size: int,
    color: (int, int, int),
    font_align="center"
):
    font = ImageFont.truetype(font_filepath, size=font_size)
    draw = ImageDraw.Draw(Image.new("RGBA", (1, 1), (0, 0, 0, 0)))
    
    bbox = draw.multiline_textbbox((0, 0), text, font=font, align=font_align)
    width = math.ceil(bbox[2] - bbox[0])
    height = math.ceil(bbox[3] - bbox[1])
    
    img = Image.new("RGBA", (1000, 500), (0, 0, 0, 0))
    draw = ImageDraw.Draw(img)
    
    if font_align == "center":
        draw_point = ((width - (bbox[2] - bbox[0])) // 2, 0)
    elif font_align == "right":
        draw_point = (width - (bbox[2] - bbox[0]), 0)
    else:
        draw_point = (0, 0)
    
    draw.multiline_text(draw_point, text, font=font, fill=color, align=font_align)
    
    return img

def add_line_breaks(text, max_characters_per_line):
    lines = []
    current_line = ""

    for word in text.split():
        if len(current_line) + len(word) + 1 <= max_characters_per_line:
            if current_line:
                current_line += " "
            current_line += word
        else:
            lines.append(current_line)
            current_line = word

    if current_line:
        lines.append(current_line)

    return "\n".join(lines)
def hex_to_rgb(hex_code):
    hex_code = hex_code.lstrip('#')

    if not all(c in '0123456789ABCDEFabcdef' for c in hex_code):
        raise ValueError("Invalid hex code format")

    r = int(hex_code[0:2], 16)
    g = int(hex_code[2:4], 16)
    b = int(hex_code[4:6], 16)

    return (b, g, r)

def create_img1(imagep, text_overlay, color):

  s3_client.download_file('inpaintingnitin', 'D1/1.jpg', '1.jpg')
  s3_client.download_file('inpaintingnitin', 'D1/2.jpg', '2.jpg')
  s3_client.download_file('inpaintingnitin', 'D1/3.jpg', '3.jpg')
  s3_client.download_file('inpaintingnitin', 'D1/Screenshot 2023-08-25 at 2.30.38 PM.png', 'img.png')
  s3_client.download_file('inpaintingnitin', 'Poppins-Medium.ttf', 'Poppins-Medium.ttf')

  image_path = 'img.png'
  image = cv2.imread(image_path)

  mask1_path = '3.jpg'
  mask = cv2.imread(mask1_path, cv2.IMREAD_GRAYSCALE)

  foreground_image_path =imagep
  foreground_image = cv2.imread(foreground_image_path)

  prim_color = hex_to_rgb(color)
  color = generate_shades((prim_color[0],prim_color[1],prim_color[2]), 10)
  new_color = color[1]

  output_image = image.copy()

  for i in range(image.shape[0]):
      for j in range(image.shape[1]):
          if mask[i, j] > 150:
              output_image[i, j] = new_color

  output_image_path = 'output_image.jpg'
  cv2.imwrite(output_image_path, output_image)

  image = output_image

  mask2_path = '2.jpg'
  mask = cv2.imread(mask2_path, cv2.IMREAD_GRAYSCALE)

  new_color = color[8]

  output_image = image.copy()

  for i in range(image.shape[0]):
      for j in range(image.shape[1]):
          if mask[i, j] > 150:
              output_image[i, j] = new_color

  output_image_path = 'output_image.jpg'
  cv2.imwrite(output_image_path, output_image)

  background_image = output_image

  mask_path = '1.jpg'
  mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

  width, height = background_image.shape[1], background_image.shape[0]
  white_image = np.ones((height, width, 3), dtype=np.uint8) * 255

  foreground_image = cv2.resize(foreground_image, (755, 890))
  x_offset = 1000
  y_offset = 50

  for y in range(foreground_image.shape[0]):
      for x in range(foreground_image.shape[1]):
          if x + x_offset < width and y + y_offset < height:
              white_image[y + y_offset, x + x_offset] = foreground_image[y, x]

  foreground_image = white_image
  mask = cv2.resize(mask, (background_image.shape[1], background_image.shape[0]))

  inv_mask = np.zeros_like(mask)

  for i in range(mask.shape[0]):
      for j in range(mask.shape[1]):
          if mask[i, j] <=150:
              inv_mask[i, j] = 255
          else:
              inv_mask[i, j] = 0

  overlay = cv2.bitwise_and(background_image, background_image, mask=inv_mask)
  overlay += cv2.bitwise_and(foreground_image, foreground_image, mask=mask)

  output_image_path = 'output_image.jpg'
  cv2.imwrite(output_image_path, overlay)

  img = text_to_image(add_line_breaks(text_overlay, 13),"Poppins-Medium.ttf",100,color[8] )
  img.save("outputtext.png", format="PNG")

  text = cv2.imread('outputtext.png',cv2.IMREAD_UNCHANGED)
  text = cv2.resize(text,(0,0),None,1,1)

  blue_channel, green_channel, red_channel, alpha_channel = cv2.split(text)

  inverted_image = cv2.merge([red_channel, green_channel, blue_channel, alpha_channel])

  final_img = cvzone.overlayPNG(overlay,inverted_image,[180,400])

  blue_channel, green_channel, red_channel = cv2.split(final_img)

  inverted_image = cv2.merge([red_channel, green_channel, blue_channel])
  return inverted_image

def create_img2(imagep, text_overlay, color):

  s3_client.download_file('inpaintingnitin', 'D2/1.jpg', '1.jpg')
  s3_client.download_file('inpaintingnitin', 'D2/2.jpg', '2.jpg')
  s3_client.download_file('inpaintingnitin', 'D2/3.jpg', '3.jpg')
  s3_client.download_file('inpaintingnitin', 'D2/Screenshot 2023-08-25 at 2.32.07 PM copy.jpg', 'img.jpg')
  s3_client.download_file('inpaintingnitin', 'Poppins-Medium.ttf', 'Poppins-Medium.ttf')


  image_path = 'img.jpg'
  image = cv2.imread(image_path)

  mask1_path = '1.jpg'
  mask = cv2.imread(mask1_path, cv2.IMREAD_GRAYSCALE)

  foreground_image_path =imagep
  foreground_image = cv2.imread(foreground_image_path)

  prim_color = hex_to_rgb(color)
  color = generate_shades((prim_color[0],prim_color[1],prim_color[2]), 10)
  new_color = color[1]

  output_image = image.copy()

  for i in range(image.shape[0]):
      for j in range(image.shape[1]):
          if mask[i, j] > 150:
              output_image[i, j] = new_color

  output_image_path = 'output_image.jpg'
  cv2.imwrite(output_image_path, output_image)

  image = output_image

  mask2_path = '2.jpg'
  mask = cv2.imread(mask2_path, cv2.IMREAD_GRAYSCALE)

  new_color = color[8]

  output_image = image.copy()

  for i in range(image.shape[0]):
      for j in range(image.shape[1]):
          if mask[i, j] > 150:
              output_image[i, j] = new_color

  output_image_path = 'output_image.jpg'
  cv2.imwrite(output_image_path, output_image)

  background_image = output_image

  mask_path = '3.jpg'
  mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

  width, height = background_image.shape[1], background_image.shape[0]
  white_image = np.ones((height, width, 3), dtype=np.uint8) * 255

  foreground_image = cv2.resize(foreground_image, (720, 745))
  x_offset = 125
  y_offset = 123

  for y in range(foreground_image.shape[0]):
      for x in range(foreground_image.shape[1]):
          if x + x_offset < width and y + y_offset < height:
              white_image[y + y_offset, x + x_offset] = foreground_image[y, x]

  foreground_image = white_image
  mask = cv2.resize(mask, (background_image.shape[1], background_image.shape[0]))

  inv_mask = np.zeros_like(mask)

  for i in range(mask.shape[0]):
      for j in range(mask.shape[1]):
          if mask[i, j] <=150:
              inv_mask[i, j] = 255
          else:
              inv_mask[i, j] = 0

  overlay = cv2.bitwise_and(background_image, background_image, mask=inv_mask)
  overlay += cv2.bitwise_and(foreground_image, foreground_image, mask=mask)

  output_image_path = 'output_image.jpg'
  cv2.imwrite(output_image_path, overlay)

  img = text_to_image(add_line_breaks(text_overlay, 13),"Poppins-Medium.ttf",100,color[1] )
  img.save("outputtext.png", format="PNG")

  text = cv2.imread('outputtext.png',cv2.IMREAD_UNCHANGED)
  text = cv2.resize(text,(0,0),None,1,1)

  blue_channel, green_channel, red_channel, alpha_channel = cv2.split(text)

  inverted_image = cv2.merge([red_channel, green_channel, blue_channel, alpha_channel])

  final_img = cvzone.overlayPNG(overlay,inverted_image,[980,400])

  blue_channel, green_channel, red_channel = cv2.split(final_img)

  inverted_image = cv2.merge([red_channel, green_channel, blue_channel])
  return inverted_image

def exec(image, text, color):
  image1 = create_img1(image, text, color)
  image2 = create_img2(image, text, color)
  return image1,image2

interface = gr.Interface(
    fn=exec,
    inputs=[gr.Image(type='filepath'),
            gr.Textbox(label = 'Text'),
            gr.ColorPicker(label="Theme Color"),],
    outputs=["image","image"],
    title="Please use square images for best results",
    allow_flagging='never',
    theme="default",
    cache_examples=False,
    ).launch( debug=True, share=False)