#!/usr/bin/env python

from __future__ import annotations

import gradio as gr

from model_list import ModelList

DESCRIPTION = '# Explore Biology & Biochemistry Foundation Models 🧬'
NOTES = '''
Thanks to the following folks who have made suggestions to this list!
- [Shelby](https://twitter.com/shelbynewsad), author of [this nice model list](https://compoundvc.notion.site/compoundvc/474885e638e94e44a1aab4d3124e3d6a?v=299bce7af785413da4c9f36837c03aaf) for some of the models listed here! 
- [Valentyn Bezshapkin](https://twitter.com/valentynbez)
- [Payel Das](https://twitter.com/payel791)
'''
FOOTER = ''''''

def main():
    model_list = ModelList()

    with gr.Blocks(css='style.css') as demo:
        gr.Markdown(DESCRIPTION)

        search_box = gr.Textbox(
            label='Search Model Name',
            placeholder=
            'You can search for titles with regular expressions. e.g. (?<!sur)face',
            max_lines=1)
        
        case_sensitive = gr.Checkbox(label='Case Sensitive')
        
        filter_names = gr.CheckboxGroup(choices=[
            'Paper',
            'Code',
            'Model Weights',
        ], label='Filter')

        data_type_names = [
            'DNA', 'scRNA', 'scRNA perturbation', 'protein language model', 'protein structure prediction', 
            'protein generation', 'protein function prediction', 'antibody structure prediction', 'antibody language model', 'molecules',
            'ligand generation', 'reaction-to-enzyme', 'enzyme generation', 'epigenomics', 'molecular docking' 
        ]

        data_types = gr.CheckboxGroup(choices=data_type_names,
                                      value=data_type_names,
                                      label='Type')

        years = ['2020', '2021', '2022', '2023']

        years_checkbox = gr.CheckboxGroup(choices=years, value=years, label='Year of Publication/Preprint')

        # model_type_names = [
        #     'GPT2', 'GPT-Neo', 'GPT-NeoX', 'ESM', 'BERT', 'RoBERTa', 'BART', 'T5', 'MPNN', 'diffusion', 'custom model'
        # ]

        # model_types = gr.CheckboxGroup(choices=model_type_names,
        #                                value=model_type_names,
        #                                label='Base Model')
        
        search_button = gr.Button('Search')

        number_of_models = gr.Textbox(label='Number of Models Found')
        table = gr.HTML(show_label=False)

        gr.Markdown(NOTES)
        gr.Markdown(FOOTER)

        demo.load(fn=model_list.render,
                  inputs=[
                      search_box,
                      case_sensitive,
                      filter_names,
                      data_types,
                      years_checkbox,
                      #model_types
                  ],
                  outputs=[
                      number_of_models,
                      table,
                  ])
        search_box.submit(fn=model_list.render,
                          inputs=[
                              search_box,
                              case_sensitive,
                              filter_names,
                              data_types,
                              years_checkbox,
                              #model_types
                          ],
                          outputs=[
                              number_of_models,
                              table,
                          ])

        search_button.click(fn=model_list.render,
                            inputs=[
                                search_box,
                                case_sensitive,
                                filter_names,
                                data_types,
                                years_checkbox,
                                #model_types
                            ],
                            outputs=[
                                number_of_models,
                                table,
                            ])

    demo.launch(enable_queue=True, share=False)


if __name__ == '__main__':
    main()