prettify and add watermark

1 files changed, 111 insertions, 27 deletions

diff --git a/main.py b/main.py
index 300691d..5f0b97f 100644
--- a/main.py
+++ b/main.py
@@ -9,7 +9,25 @@ import sv_ttk
 
 import os
 
-def pyramid_list(lst):
+canadian_companies = []
+
+def pyramid_list(lst, sort_canadian=True):
+    old_lst = lst
+
+    top = []
+    bottom = []
+
+    for idx, company in enumerate(lst):
+        if idx==0:
+            top.append(company)
+        else:
+            if company in canadian_companies:
+                bottom.append(company)
+            else:
+                top.append(company)
+
+    lst = top + bottom
+
     i = 0
     rows = 0
     pyramid = []
@@ -28,7 +46,40 @@ def pyramid_list(lst):
                 row.append(lst[i])
                 i += 1
         pyramid.append(row)
+
+    if len(pyramid) >= 2:
+        if len(pyramid[-1]) < len(pyramid[-2]):
+            pyramid[-2].extend(pyramid[-1])
+            pyramid.pop()
+
+    print("Checking for canadian companies")
+    print(canadian_companies)
+    if len(pyramid) > 2:
+        print("Pyramid has more than two rows, checking to ensure Canadian companies are at the bottom")
+        print(len(pyramid[-1]) - len(pyramid[-2]))
+        for row in pyramid:
+            print(len(row))
+            print(row)
+        if ((len(pyramid[-1]) - len(pyramid[-2])) > 1):
+            print("Last two rows can be balanced if needed")
+            for _ in range((len(pyramid[-1]) - len(pyramid[-2]))-1):
+                company_to_add = None
+                company_idx = 0
+                for idx, company in enumerate(pyramid[-1]):
+                    print(f'Checking {company}')
+                    if company not in canadian_companies:
+                        company_to_add = company
+                        company_idx = idx
+                        print(f'Shifting {company_to_add}')
+                if company_to_add is not None:
+                    pyramid[-1].pop(company_idx)
+                    pyramid[-2].append(company_to_add)
     
+    if len(pyramid) > 3:
+        if len(pyramid[-1]) < len(pyramid[-2]):
+            pyramid[-3].append(pyramid[-2][-1])
+            pyramid[-2].pop(-1)
+
     return pyramid
 
 class CompanySelector:
@@ -92,10 +143,12 @@ class CompanySelector:
         company_scores = {}
 
         # For each 'Company Name', company_scores[company] = 'Weighting' for that company
-        for company, ticker, weighting, dividend in zip(selected_df["Company Name"], selected_df["Symbol"], selected_df["Weighting"], selected_df["Dividend?"]):
+        for company, ticker, weighting, dividend, currency in zip(selected_df["Company Name"], selected_df["Symbol"], selected_df["Weighting"], selected_df["Dividend?"], selected_df["Currency"]):
             extra_char = ""
             if dividend.strip() == "Y":
                 extra_char = "*"
+            if currency == "CAD":
+                canadian_companies.append(f'{company} ({ticker}){extra_char}')
             company_scores[f'{company} ({ticker}){extra_char}'] = weighting
 
         from PIL import Image, ImageDraw, ImageFont
@@ -111,12 +164,18 @@ class CompanySelector:
         # Initialize some parameters
         img_width = 3300
         img_height = 1550 #2550
-        block_color = (0, 0, 255)  # Blue color
-        dividend_color = (0, 255, 0)  # Green color
+        overall_height = 2550
+        block_color = (3, 37, 126)  # Blue color
+        dividend_color = block_color #(1,50,32)  # Green color
         font_color = (255, 255, 255)  # White color
+        normal_block = (3, 37, 126) # Dark Blue  
         padding = 10  # Padding around blocks
         radius = 20 
+        logo_padding = 100 # top padding for logo
+        pyramid_padding_bottom = 350 # bottom padding for pyramid
         max_companies_in_row = 5
+        primary_font = "./assets/baskerville.ttf"
+        secondary_font = "./assets/gill_sans_bold.ttf"
 
         # Group the companies by their scores
         pyramid = {}
@@ -157,11 +216,6 @@ class CompanySelector:
 
         companies = pyramid_list(companies)
 
-        if len(companies) >= 2:
-            if len(companies[-1]) < len(companies[-2]):
-                companies[-2].extend(companies[-1])
-                companies.pop()
-
         for company_row in companies:
             print(company_row)
 
@@ -228,17 +282,17 @@ class CompanySelector:
                     block_color = dividend_color
                     num_dividends += 1
                 else:
-                    block_color = (0, 0, 255)
+                    block_color = normal_block
 
                 # Draw the block
                 d.rounded_rectangle([x, y, x + block_width, y + block_height], fill=block_color, radius=radius)
 
                 # Adjust font size based on the length of the company name and block size
                 font_size = min_font_size #min(block_width // (len(company) // 2 + 1), block_height // 2)
-                fnt = ImageFont.truetype('./assets/arial.ttf', font_size)
+                fnt = ImageFont.truetype(secondary_font, font_size)
 
                 # Implement word wrap for the company name
-                wrapped_company = wrap_text(company, block_width // font_size)
+                wrapped_company = wrap_text(company, (block_width // font_size)*1.5)
 
                 # Draw the company name
                 bbox = d.textbbox((x, y), wrapped_company, font=fnt)
@@ -246,7 +300,7 @@ class CompanySelector:
                 text_height = bbox[3] - bbox[1]
                 text_x = x + (block_width - text_width) // 2
                 text_y = y + (block_height - text_height) // 2
-                d.text((text_x, text_y), wrapped_company, font=fnt, fill=font_color)
+                d.text((text_x, text_y), wrapped_company, font=fnt, fill=font_color, align='center', spacing=10)
 
 
         print(image_file_path)
@@ -264,39 +318,69 @@ class CompanySelector:
 
         new_img.paste(image_logo, (left, top))
 
-        final_img = Image.new("RGB", (img_width, 2550), "white")
-        final_img.paste(new_img, (0, 0))
+        final_img = Image.new("RGB", (img_width, overall_height), "white")
+        final_img.paste(new_img, (0, logo_padding))
 
         text1 = self.pyramid_title_string.get()
         text2 = f"({num_dividends} Dividend payors - all identified by asterisk)"
+        text3 = "FOR INTERNAL USE ONLY"
 
         draw = ImageDraw.Draw(final_img)
         font_size = 60
 
-        font1 = ImageFont.truetype('./assets/arial.ttf', font_size)
-        text_width1, _ = draw.textsize(text1, font=font1)
+        font1 = ImageFont.truetype(primary_font, font_size)
+        text_width1 = draw.textlength(text1, font=font1)
         while text_width1 > final_img.width:
             font_size -= 1
-            font1 = ImageFont.truetype('./assets/arial.ttf', font_size)
-            text_width1, _ = draw.textsize(text1, font=font1)
+            font1 = ImageFont.truetype(primary_font, font_size)
+            text_width1 = draw.textlength(text1, font=font1)
 
-        font2 = ImageFont.truetype('./assets/arial.ttf', font_size)
-        text_width2, _ = draw.textsize(text2, font=font2)
+        font2 = ImageFont.truetype(primary_font, font_size)
+        text_width2 = draw.textlength(text2, font=font2)
         while text_width2 > final_img.width:
             font_size -= 1
-            font2 = ImageFont.truetype('./assets/arial.ttf', font_size)
-            text_width2, _ = draw.textsize(text2, font=font2)
+            font2 = ImageFont.truetype(primary_font, font_size)
+            text_width2 = draw.textlength(text2, font=font2)
+
+        new_font_size = 30
+        font3 = ImageFont.truetype(secondary_font, new_font_size)
+        _,t,_,b = draw.textbbox((100,100), text3, font=font3)
+        while (b-t) < (0.8*pyramid_padding_bottom):
+            new_font_size += 1
+            font3 = ImageFont.truetype(secondary_font, new_font_size)
+            _,t,_,b = draw.textbbox((100,100), text3, font=font3)
+        text_width3 = draw.textlength(text3, font=font3)
+        while text_width3 > (final_img.width*0.8):
+            new_font_size -= 1
+            font3 = ImageFont.truetype(secondary_font, new_font_size)
+            text_width3 = draw.textlength(text3, font=font3)
+        _,t,_,b = draw.textbbox((100,100),text3,font=font3)
 
         start_x1 = (final_img.width - text_width1) // 2
-        start_y1 = image_logo.height + (text_height // 3)
+        start_y1 = image_logo.height + logo_padding + 50
 
         start_x2 = (final_img.width - text_width2) // 2
-        start_y2 = image_logo.height + 4 * (text_height // 3)
+        start_y2 = image_logo.height + logo_padding + 125
+
+        start_x3 = (final_img.width - text_width3) // 2
+        start_y3 = (final_img.height - pyramid_padding_bottom + ((pyramid_padding_bottom)-(b-t))//2)
         
         draw.text((start_x1, start_y1), text1, font=font1, fill="black")
         draw.text((start_x2, start_y2), text2, font=font2, fill="black")
-
-        final_img.paste(img, (0, l_height + (2550 - l_height - img_height)))
+        draw.text((start_x3, start_y3), text3, font=font3, fill=(105,105,105))
+
+        final_img.paste(img, (0, l_height + (overall_height - l_height - img_height) - pyramid_padding_bottom) )
+
+        # draw = ImageDraw.Draw(final_img)
+        # watermark_text = "FOR INTERNAL USE ONLY"
+        # watermark_font = ImageFont.truetype(secondary_font, 72)
+        # watermark_width = draw.textlength(watermark_text, watermark_font)
+        # draw.text(
+        #     ((final_img.width - watermark_width)//2,overall_height-pyramid_padding_bottom),
+        #     watermark_text,
+        #     font=watermark_font,
+        #     align='center'
+        #     )
 
         final_img.save(image_file_path)