from metaflow import FlowSpec, Parameter, conda, conda_base, schedule, JSONType

from common.constants import SENTENCE_TRANSFORMERS_S3_PREFIX, IMAGE_W_NVIDIA_DRIVER
from common.metaflow_utils import step, get_environ
from common.utils import partition_list

BATCH_SIZE = 2000
MODEL_S3_PREFIX = f'{SENTENCE_TRANSFORMERS_S3_PREFIX}/stsb-bert-base/'


@schedule(hourly=True)
@conda_base(python='3.8', libraries={'snowflake-connector-python': '2.3.8'})
class RelationshipFeatureResolution(FlowSpec):
    metaflow_environ = Parameter('metaflow_environ', type=JSONType, default=get_environ)

    @step
    def start(self):
        from sentence_transformer.feature.relationship.store import create_table, query_content
        from sentence_transformer.feature.store import create_schema
        create_schema(self.ctx.secret_id)
        create_table(self.ctx.secret_id)
        self.batches = partition_list(query_content(self.ctx.secret_id, batch_size=BATCH_SIZE), batch_size=200)
        self.next(self.resolve_relationships, foreach='batches')

    @conda(libraries={
        'pandas': '1.1.3', 'numpy': '1.20.1', 'scikit-learn': '0.24.1',
        'sentence-transformers': '0.4.1', 'gensim': '3.8.3'
    })
    @step
    def resolve_relationships(self):
        from sentence_transformer.feature.relationship.relationship import get_initial
        from sentence_transformer.feature.relationship.store import save
        relationships = get_initial(
            self.input, MODEL_S3_PREFIX, self.ctx.environment, self.ctx.log)
        <http://self.ctx.log.info|self.ctx.log.info>(f'saving {len(relationships)} relationship descriptions')
        success, num_chunks, num_rows, output = save(relationships, self.ctx.secret_id)
        <http://self.ctx.log.info|self.ctx.log.info>(
            f'save completed with success: {success}, num_chunks: {num_chunks}, num_rows: {num_rows}, output: {output}')
        self.next(self.join)

    @conda(disabled=True)
    @step
    def join(self, inputs):
        self.next(self.end)

    @conda(disabled=True)
    @step
    def end(self):
        pass


if __name__ == '__main__':
    RelationshipFeatureResolution()