from imgofai import *

import matplotlib.pyplot as plt

import PIL

import numpy as np
from pathlib import Path

from imgofai.tree import img2df, df2xy
import pandas as pd 
from sklearn.tree import DecisionTreeRegressor
from sklearn.preprocessing import FunctionTransformer
from sklearn.pipeline import make_pipeline

import streamlit as st

import requests

st.write("# Images of AI Demo")

st.write("This page demonstrates how I created the images I submitted for [Better Images of AI project](https://betterimagesofai.org/images)")

def add_radial_features(X,y=None):
    assert isinstance(X, pd.DataFrame), "X is not a dataframe"
    xp = X.copy()
    xp['dim0'] = np.sqrt(((X - X.mean())**2).sum(axis=1))
    xp['dim1'] = np.arctan2(X['dim1'],X['dim0'])

    X = pd.concat([
        X,
        xp,
        ],axis=1)

    return X



def make_tree_approximator(radial = False, max_depth=4):
    if radial:
        model = make_pipeline(
            FunctionTransformer(add_radial_features),
            DecisionTreeRegressor(max_depth=max_depth),
        )
    else:
        model =  DecisionTreeRegressor(max_depth=max_depth)


    model.fit(x_raw, y)
    pred = PIL.Image.fromarray(
        model.predict(x_raw).reshape(img_array.shape).round().astype("uint8")
    )
    score = model.score(x_raw, y)
    return pred


st.write("## Try it out yourself:")

url = st.text_input("Image url (make sure it ends with .png/.jpg/...):", "https://images.unsplash.com/reserve/bOvf94dPRxWu0u3QsPjF_tree.jpg?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=1752&q=80")
img = PIL.Image.open(requests.get(url, stream=True).raw)

img_array = np.array(img)
df = img2df(img_array)
x_raw, y = df2xy(df)



ccol1, ccol2, _ = st.columns(3)

with ccol1:
    max_depth1 = st.slider("max depth left:",1,12,2)
    radial1 = st.checkbox("radial features left", value=False)

with ccol2:
    max_depth2 = st.slider("max depth middle:",1,12,6)
    radial2 = st.checkbox("radial features middle", value=False)


st.write("## Output")
col1, col2, col3 = st.columns(3)

with col1:
    left_img = make_tree_approximator(radial1, max_depth=max_depth1)
    st.image(left_img)

with col2:
    mid_img = make_tree_approximator(radial2, max_depth=max_depth2)
    st.image(mid_img)

with col3:
    st.image(img)