import json
import pytest
import requests_mock

from transfer.predict_help import (
    convert_dict_to_actions,
    parse_item_page_amz, parse_results_amz,
    parse_item_page_ebay, parse_results_ebay,
    parse_item_page_ws, parse_results_ws,
    Page, WEBSHOP_URL, WEBSHOP_SESSION
)

@requests_mock.Mocker(kw="mock")
def test_parse_item_page_ws(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_item_page_ws", "rb")
    mock_body = mock_file.read()
    mock_file.close()

    mock_desc_file = open("tests/transfer/mocks/mock_parse_item_page_ws_desc", "rb")
    mock_desc_body = mock_desc_file.read()
    mock_desc_file.close()

    mock_feat_file = open("tests/transfer/mocks/mock_parse_item_page_ws_feat", "rb")
    mock_feat_body = mock_feat_file.read()
    mock_feat_file.close()

    mock_asin = "B09P87V3LZ"
    mock_query = "red basketball shoes"
    mock_options = {}

    # Invoke function, check response
    query_str = '+'.join(mock_query.split())
    options_str = json.dumps(mock_options)
    url = (
        f"{WEBSHOP_URL}/item_page/{WEBSHOP_SESSION}/"
        f"{mock_asin}/{query_str}/1/{options_str}"
    )
    url_desc = (
        f"{WEBSHOP_URL}/item_sub_page/{WEBSHOP_SESSION}/"
        f"{mock_asin}/{query_str}/1/Description/{options_str}"
    )
    url_feat = (
        f"{WEBSHOP_URL}/item_sub_page/{WEBSHOP_SESSION}/"
        f"{mock_asin}/{query_str}/1/Features/{options_str}"
    )
    print(f"Item Page URL: {url}")
    print(f"Item Description URL: {url_desc}")
    print(f"Item Features URL: {url_feat}")

    kwargs["mock"].get(url, content=mock_body)
    kwargs["mock"].get(url_desc, content=mock_desc_body)
    kwargs["mock"].get(url_feat, content=mock_feat_body)

    output = parse_item_page_ws(mock_asin, mock_query, 1, mock_options)
    expected = {
        'MainImage': 'https://m.media-amazon.com/images/I/51ltvkzGhGL.jpg',
        'Price': '100.0',
        'Rating': 'N.A.',
        'Title': 'PMUYBHF Womens Fashion Flat Shoes Comfortable Running Shoes '
            'Sneakers Tennis Athletic Shoe Casual Walking Shoes',
        'asin': mock_asin,
        'option_to_image': {
            '6.5': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27size%27:%20%276.5%27%7D',
            '7.5': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27size%27:%20%277.5%27%7D',
            '8': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27size%27:%20%278%27%7D',
            '8.5': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27size%27:%20%278.5%27%7D',
            '9': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27size%27:%20%279%27%7D',
            'black': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27color%27:%20%27black%27%7D',
            'purple': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27color%27:%20%27purple%27%7D',
            'red': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27color%27:%20%27red%27%7D'
        },
        'options': {
            'color': ['black', 'purple', 'red'],
            'size': ['6.5', '7.5', '8', '8.5', '9']
        },
        'BulletPoints': 'Pure Running Shoe\nComfort Flat Sneakers\n[FEATURES]: Soles with unique non-slip pattern, it has great '
                'abrasion resistant and provide protection when you walking or running. (Pure Running Shoe Mesh Walking Shoes Fashion '
                'Sneakers Slip On Sneakers Wedge Platform Loafers Modern Walking Shoes Sock Sneakers Platform Loafers Shoes Non Slip '
                'Running Shoes Athletic Tennis Shoes Blade Type Sneakers Lace-up Sneaker) sole\n[WIDE ANKLE DESIGN]: Perfect accord with human body '
                'engineering, green, healthy concept design make the walking shoes wear more comfortable, wide width wlking shoes. (Low '
                'Top Walking Shoes Fashion Canvas Sneakers Slip On Shoes Casual Walking Shoes Hidden Wedge Sneaker Low Top Canvas '
                'Sneakers Lace-up Classic Casual Shoes Walking Tennis Shoes Lightweight Casual Sneakers Slip on Sock Sneakers Air '
                'Cushion Platform Loafers Slip-On Mule Sneaker )\n[CUSHION WITH ARCH SUPPORT]: Gives you a comfort for all day '
                'long. Wear these lightweight walking shoes, let every step of moving on a comfortable feeling. (Fashion Casual Shoes '
                'Athletic Workout Shoes Fitness Sneaker Athletic Running Shoes Air Cushion Sneakers Stylish Athletic Shoes Lace Up '
                'Canvas Shoes Slip on Walking Shoe Fashion Sneakers Low Top Classic Sneakers Comfort Fall Shoes Memory Foam Slip On '
                'Sneakers Air Cushion Sneakers Running Walking Shoes)\n[NON-SLIP SOLE]: Made from ultra soft and lightweight RUBBER '
                'material,with the function of shock absorbing and cushioning,offering the best durability and traction. (Wedge '
                'Sneakers Walking Tennis Shoes Slip On Running Shoes Lightweight Fashion Sneakers Fashion Travel Shoes Walking '
                'Running Shoes Non Slip Running Shoes Athletic Tennis Sneakers Sports Walking Shoes Platform Fashion Sneaker '
                'Memory Foam Tennis Sneakers Running Jogging Shoes Sock Sneakers Canvas Fashion Sneakers)\n[OCCASIONS]: Ultra lightweight design provides actual '
                'feelings of being barefooted and like walking on the feather, perfect for walking, hiking, bike riding, working, '
                'shopping, indoor, outdoor, casual, sports, travel, exercise, vacation, and etc. (Flat Fashion Sneakers Lightweight '
                'Walking Sneakers Platform Loafers Sport Running Shoes Casual Flat Loafers Slip-On Sneaker Casual Walking Shoes High Top '
                'Canvas Sneakers Lace up Sneakers Workout Walking Shoes Tennis Fitness Sneaker)\n'
                '[Customers Are Our Priority]: We follow the principle of customer first, so if you encounter any problems after '
                'buying shoes, we will try our best to solve them for you. (Breathable Air Cushion Sneakers Walking Tennis Shoes Air '
                'Athletic Running Shoes Air Cushion Shoes Mesh Sneakers Fashion Tennis Shoes Jogging Walking Sneakers Breathable '
                'Casual Sneakers Fashion Walking Shoes Athletic Running Sneakers Walking Work Shoes Air Running Shoes Slip on '
                'Sneakers Mesh Walking Shoes)',
        'Description': 'Here Are The Things You Want To Knowa─=≡Σ(((つ̀ώ)つSTORE INTRODUCTION:>>>>Our store helps our customers improve their '
               'quality of life~As a distributor, we value quality and service.Focus on the high quality and durability of the '
               'product.Committed to creating a store that satisfies and reassures our customers.TIPS:>>>>1. Please allow minor errors '
               'in the data due to manual measurements.2. Due to the color settings of the display, the actual color may be slightly '
               'different from the online image.QUALITY PROMISE:>>>>Our goal is to continuously provide a range of quality products.We '
               'place a huge emphasis on the values of quality and reliability.We have always insisted on fulfilling this '
               'commitment.In short, we want our customers to have the same great product experience every time and be trusted to deliver '
               'on this commitment.Please give us a chance to serve you.OTHER:>>>>athletic sneaker laces athletic sneakers white '
               'athletic sneakers for women clearance leather Sneaker leather sneakers women leather sneakers for menleather sneaker laces '
               'leather sneaker platform basketball shoes basketball shoes for men basketball shoe laces basketball shoe grip basketball '
               'shoes for women fitness shoes for men fitness shoes women workout fitness shoes women fitness shoes women size 5 '
               'fitness shoes men workout fitness shoes for men high top sneakers for women walking shoes sneakers with arch support for women'
    }
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_item_page_ebay(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_item_page_ebay", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_asin = "403760625150"

    # Invoke function, check response
    kwargs["mock"].get(f"https://www.ebay.com/itm/{mock_asin}", content=mock_body)
    output = parse_item_page_ebay(mock_asin)
    expected = {
        'BulletPoints': 'Item specifics Condition:New without box: A brand-new, '
            'unused, and unworn item (including handmade items) that is '
            'not in ...  Read moreabout the conditionNew without box: A '
            'brand-new, unused, and unworn item (including handmade '
            'items) that is not in original packaging or may be missing '
            'original packaging materials (such as the original box or '
            'bag). The original tags may not be attached. For example, '


                'buying shoes, we will try our best to solve them for you. (Breathable Air Cushion Sneakers Walking Tennis Shoes Air '
                'Athletic Running Shoes Air Cushion Shoes Mesh Sneakers Fashion Tennis Shoes Jogging Walking Sneakers Breathable '
                'Casual Sneakers Fashion Walking Shoes Athletic Running Sneakers Walking Work Shoes Air Running Shoes Slip on '
                'Sneakers Mesh Walking Shoes)',
        'Description': 'Here Are The Things You Want To Knowa─=≡Σ(((つ̀ώ)つSTORE INTRODUCTION:>>>>Our store helps our customers improve their '
               'quality of life~As a distributor, we value quality and service.Focus on the high quality and durability of the '
               'product.Committed to creating a store that satisfies and reassures our customers.TIPS:>>>>1. Please allow minor errors '
               'in the data due to manual measurements.2. Due to the color settings of the display, the actual color may be slightly '
               'different from the online image.QUALITY PROMISE:>>>>Our goal is to continuously provide a range of quality products.We '
               'place a huge emphasis on the values of quality and reliability.We have always insisted on fulfilling this '
               'commitment.In short, we want our customers to have the same great product experience every time and be trusted to deliver '
               'on this commitment.Please give us a chance to serve you.OTHER:>>>>athletic sneaker laces athletic sneakers white '
               'athletic sneakers for women clearance leather Sneaker leather sneakers women leather sneakers for menleather sneaker laces '
               'leather sneaker platform basketball shoes basketball shoes for men basketball shoe laces basketball shoe grip basketball '
               'shoes for women fitness shoes for men fitness shoes women workout fitness shoes women fitness shoes women size 5 '
               'fitness shoes men workout fitness shoes for men high top sneakers for women walking shoes sneakers with arch support for women'
    }
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_item_page_ebay(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_item_page_ebay", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_asin = "403760625150"

    # Invoke function, check response
    kwargs["mock"].get(f"https://www.ebay.com/itm/{mock_asin}", content=mock_body)
    output = parse_item_page_ebay(mock_asin)
    expected = {
        'BulletPoints': 'Item specifics Condition:New without box: A brand-new, '
            'unused, and unworn item (including handmade items) that is '
            'not in ...  Read moreabout the conditionNew without box: A '
            'brand-new, unused, and unworn item (including handmade '
            'items) that is not in original packaging or may be missing '
            'original packaging materials (such as the original box or '
            'bag). The original tags may not be attached. For example, '
            'new shoes (with absolutely no signs of wear) that are no '
            'longer in their original box fall into this category.  See '
            'all condition definitionsopens in a new window or tab  '
            'Closure:Lace Up US Shoe Size:10 Occasion:Activewear, Casual '
            'Silhouette:Puma Fabric Type:Mesh Vintage:No Cushioning '
            'Level:Moderate Department:Men Style:Sneaker Outsole '
            'Material:Rubber Features:Breathable, Comfort, Cushioned, '
            'Performance Season:Fall, Spring, Summer, Winter '
            'Idset_Mpn:193990-21 Shoe Shaft Style:Low Top Style '
            'Code:193990-16 Pattern:Solid Character:J. Cole Lining '
            'Material:Synthetic Color:Red Brand:PUMA Type:Athletic '
            'Customized:No Model:RS-Dreamer Theme:Sports Shoe '
            'Width:Standard Upper Material:Textile Insole '
            'Material:Synthetic Performance/Activity:Basketball Product '
            'Line:Puma Dreamer',
        'Description': 'N/A',
        'MainImage': 'https://i.ebayimg.com/images/g/4ggAAOSwpk1ioTWz/s-l500.jpg',
        'Price': 'N/A',
        'Rating': None,
        'Title': "Puma RS-Dreamer J. Cole Basketball Shoes Red 193990-16 Men's Size 10.0",
        'asin': '403760625150',
        'option_to_image': {},
        'options': {},
    }
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_item_page_amz(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_item_page_amz", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_asin = "B073WRF565"

    # Invoke function, check response
    kwargs["mock"].get(f"https://www.amazon.com/dp/{mock_asin}", content=mock_body)
    output = parse_item_page_amz(mock_asin)
    expected = {
        'asin': 'B073WRF565',
        'Title': 'Amazon Basics Foldable 14" Black Metal Platform Bed Frame with Tool-Free Assembly No Box Spring Needed - Full',
        'Price': 'N/A',
        'Rating': '4.8 out of 5 stars',
        'BulletPoints': ' \n About this item    '
            'Product dimensions: 75" L x 54" W x 14" H | Weight: 41.4 pounds    '
            'Designed for sleepers up to 250 pounds    Full size platform bed frame offers a quiet, noise-free, '


            'Performance Season:Fall, Spring, Summer, Winter '
            'Idset_Mpn:193990-21 Shoe Shaft Style:Low Top Style '
            'Code:193990-16 Pattern:Solid Character:J. Cole Lining '
            'Material:Synthetic Color:Red Brand:PUMA Type:Athletic '
            'Customized:No Model:RS-Dreamer Theme:Sports Shoe '
            'Width:Standard Upper Material:Textile Insole '
            'Material:Synthetic Performance/Activity:Basketball Product '
            'Line:Puma Dreamer',
        'Description': 'N/A',
        'MainImage': 'https://i.ebayimg.com/images/g/4ggAAOSwpk1ioTWz/s-l500.jpg',
        'Price': 'N/A',
        'Rating': None,
        'Title': "Puma RS-Dreamer J. Cole Basketball Shoes Red 193990-16 Men's Size 10.0",
        'asin': '403760625150',
        'option_to_image': {},
        'options': {},
    }
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_item_page_amz(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_item_page_amz", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_asin = "B073WRF565"

    # Invoke function, check response
    kwargs["mock"].get(f"https://www.amazon.com/dp/{mock_asin}", content=mock_body)
    output = parse_item_page_amz(mock_asin)
    expected = {
        'asin': 'B073WRF565',
        'Title': 'Amazon Basics Foldable 14" Black Metal Platform Bed Frame with Tool-Free Assembly No Box Spring Needed - Full',
        'Price': 'N/A',
        'Rating': '4.8 out of 5 stars',
        'BulletPoints': ' \n About this item    '
            'Product dimensions: 75" L x 54" W x 14" H | Weight: 41.4 pounds    '
            'Designed for sleepers up to 250 pounds    Full size platform bed frame offers a quiet, noise-free, '
            'supportive foundation for a mattress. No box spring needed    Folding mechanism makes the frame easy '
            'to store and move in tight spaces    Provides extra under-the-bed storage space with a vertical clea'
            'rance of about 13 inches    \n › See more product details ',
        'Description': 'Amazon Basics Foldable, 14" Black Metal Platform Bed Frame with Tool-Free Assembly, No Box Spring Needed - Full   Amazon Basics',
        'MainImage': 'https://images-na.ssl-images-amazon.com/images/I/41WIGwt-asL.__AC_SY300_SX300_QL70_FMwebp_.jpg',
        'options': {'size': ['Twin', 'Full', 'Queen', 'King'],
        'style': ['14-Inch', '18-Inch']},
        'option_to_image': {}
    }
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_results_ebay(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_results_ebay", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_query = "red basketball shoes"

    # Invoke function, check response
    query = mock_query.replace(" ", "+")
    kwargs["mock"].get(f'https://www.ebay.com/sch/i.html?_nkw={query}&_pgn=1', content=mock_body)
    output = parse_results_ebay(mock_query, 1)
    expected = [{
		'Price': ['100.00', '150.00'],
		'Title': "Reebok Answer IV Men's Basketball Shoes",
		'asin': '175065123030'
	}, {
		'Price': '$119.90',
		'Title': "Air Jordan Stay Loyal Shoes Black Red White DB2884-001 Men's Multi "


    expected = {
        'asin': 'B073WRF565',
        'Title': 'Amazon Basics Foldable 14" Black Metal Platform Bed Frame with Tool-Free Assembly No Box Spring Needed - Full',
        'Price': 'N/A',
        'Rating': '4.8 out of 5 stars',
        'BulletPoints': ' \n About this item    '
            'Product dimensions: 75" L x 54" W x 14" H | Weight: 41.4 pounds    '
            'Designed for sleepers up to 250 pounds    Full size platform bed frame offers a quiet, noise-free, '
            'supportive foundation for a mattress. No box spring needed    Folding mechanism makes the frame easy '
            'to store and move in tight spaces    Provides extra under-the-bed storage space with a vertical clea'
            'rance of about 13 inches    \n › See more product details ',
        'Description': 'Amazon Basics Foldable, 14" Black Metal Platform Bed Frame with Tool-Free Assembly, No Box Spring Needed - Full   Amazon Basics',
        'MainImage': 'https://images-na.ssl-images-amazon.com/images/I/41WIGwt-asL.__AC_SY300_SX300_QL70_FMwebp_.jpg',
        'options': {'size': ['Twin', 'Full', 'Queen', 'King'],
        'style': ['14-Inch', '18-Inch']},
        'option_to_image': {}
    }
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_results_ebay(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_results_ebay", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_query = "red basketball shoes"

    # Invoke function, check response
    query = mock_query.replace(" ", "+")
    kwargs["mock"].get(f'https://www.ebay.com/sch/i.html?_nkw={query}&_pgn=1', content=mock_body)
    output = parse_results_ebay(mock_query, 1)
    expected = [{
		'Price': ['100.00', '150.00'],
		'Title': "Reebok Answer IV Men's Basketball Shoes",
		'asin': '175065123030'
	}, {
		'Price': '$119.90',
		'Title': "Air Jordan Stay Loyal Shoes Black Red White DB2884-001 Men's Multi "
		'Size NEW',
		'asin': '265672133690'
	}, {
		'Price': '$100.00',
		'Title': "Fila Men's Stackhouse Spaghetti Basketball Shoes Black Red White "
		'1BM01788-113',
		'asin': '175282509234'
	}, {
		'Price': ['61.99',
			'85.99'
		],
		'Title': 'Puma Disc Rebirth 19481203 Mens Black Red Synthetic Athletic '
		'Basketball Shoes',
		'asin': '313944854658'
	}, {
		'Price': '$0.01',
		'Title': "Puma RS-Dreamer J. Cole Basketball Shoes Red 193990-16 Men's Size "
		'10.0',
		'asin': '403760625150'
	}, {
		'Price': '$45.00',
		'Title': 'Nike Mens 9.5 PG 5  Maroon Red White Basketball Shoes Sneaker DM '
		'5045–601￼ Flaw',
		'asin': '115456853186'
	}, {
		'Price': ['114.90',
			'119.90'
		],
		'Title': "Air Jordan Stay Loyal Shoes White Black Red DB2884-106 Men's Multi "
		'Size NEW',
		'asin': '155046831159'
	}, {
		'Price': '$8.99',
		'Title': "Harden Volume 3 Men's Basketball Shoes Size 9.5",
		'asin': '175342407862'
	}, {
		'Price': '$59.97',
		'Title': "Men's Nike Precision 5 Basketball Shoes Gym Red Black Grey Bred "
		'Multi Size NEW',
		'asin': '134149634710'
	}]
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_results_amz(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_results_amz", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_query = "red basketball shoes"

    # Invoke function, check response
    query = mock_query.replace(" ", "+")
    kwargs["mock"].get(f"https://www.amazon.com/s?k={query}&page=1", content=mock_body)
    output = parse_results_amz(mock_query, 1)
    expected = [{
		'Price': '59.49',
		'Title': 'High Top Mens Basketball Shoes Lou Williams Streetball Master '
			'Breathable Non Slip Outdoor Sneakers Cushioning Workout Shoes for '
			'Fitness',
		'asin': 'B083QCWF61'
	}, {


	}, {
		'Price': ['114.90',
			'119.90'
		],
		'Title': "Air Jordan Stay Loyal Shoes White Black Red DB2884-106 Men's Multi "
		'Size NEW',
		'asin': '155046831159'
	}, {
		'Price': '$8.99',
		'Title': "Harden Volume 3 Men's Basketball Shoes Size 9.5",
		'asin': '175342407862'
	}, {
		'Price': '$59.97',
		'Title': "Men's Nike Precision 5 Basketball Shoes Gym Red Black Grey Bred "
		'Multi Size NEW',
		'asin': '134149634710'
	}]
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_results_amz(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_results_amz", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_query = "red basketball shoes"

    # Invoke function, check response
    query = mock_query.replace(" ", "+")
    kwargs["mock"].get(f"https://www.amazon.com/s?k={query}&page=1", content=mock_body)
    output = parse_results_amz(mock_query, 1)
    expected = [{
		'Price': '59.49',
		'Title': 'High Top Mens Basketball Shoes Lou Williams Streetball Master '
			'Breathable Non Slip Outdoor Sneakers Cushioning Workout Shoes for '
			'Fitness',
		'asin': 'B083QCWF61'
	}, {
		'Price': '45.99',
		'Title': 'Kids Basketball Shoes High-top Sports Shoes Sneakers Durable '
		'Lace-up Non-Slip Running Shoes Secure for Little Kids Big Kids and '
		'Boys Girls',
		'asin': 'B08FWWWQ11'
	}, {
		'Price': '64.99',
		'Title': 'Unisex-Adult Lockdown 5 Basketball Shoe',
		'asin': 'B0817BFNC4'
	}, {
		'Price': '63.75',
		'Title': 'Unisex-Child Team Hustle D 9 (Gs) Sneaker',
		'asin': 'B07HHTS79M'
	}, {
		'Price': '74.64',
		'Title': 'Unisex-Adult D.O.N. Issue 3 Basketball Shoe',
		'asin': 'B08N8DQLS2'
	}, {
		'Price': '104.90',
		'Title': "Men's Lebron Witness IV Basketball Shoes",
		'asin': 'B07TKMMHVB'
	}, {
		'Price': '36.68',
		'Title': "Unisex-Child Pre-School Jet '21 Basketball Shoe",
		'asin': 'B08N6VRHV4'
	}, {
		'Price': '59.98',
		'Title': "Men's Triple Basketball Shoe",
		'asin': 'B08QCL8VKM'
	}, {
		'Price': '45.98',
		'Title': 'Unisex-Child Pre School Lockdown 4 Basketball Shoe',
		'asin': 'B07HKP12DH'
	}, {
		'Price': '143.72',
		'Title': "Men's Basketball Shoes",
		'asin': 'B07SNR7HRF'
	}]
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_results_ws(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_results_ws", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_query = "red basketball shoes"

    # Invoke function, check response
    query_str = mock_query.replace(" ", "+")
    url = (
        f'{WEBSHOP_URL}/search_results/{WEBSHOP_SESSION}/'
        f'{query_str}/1'
    )
    kwargs["mock"].get(url, content=mock_body)
    output = parse_results_ws(mock_query, 1)
    expected = [{
        'Price': [24.49, 39.99],
        'Title': "BinGoDug Men's Basketball Shoes, Men's Fashion Sneakers, Air "


	}, {
		'Price': '36.68',
		'Title': "Unisex-Child Pre-School Jet '21 Basketball Shoe",
		'asin': 'B08N6VRHV4'
	}, {
		'Price': '59.98',
		'Title': "Men's Triple Basketball Shoe",
		'asin': 'B08QCL8VKM'
	}, {
		'Price': '45.98',
		'Title': 'Unisex-Child Pre School Lockdown 4 Basketball Shoe',
		'asin': 'B07HKP12DH'
	}, {
		'Price': '143.72',
		'Title': "Men's Basketball Shoes",
		'asin': 'B07SNR7HRF'
	}]
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_results_ws(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_results_ws", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_query = "red basketball shoes"

    # Invoke function, check response
    query_str = mock_query.replace(" ", "+")
    url = (
        f'{WEBSHOP_URL}/search_results/{WEBSHOP_SESSION}/'
        f'{query_str}/1'
    )
    kwargs["mock"].get(url, content=mock_body)
    output = parse_results_ws(mock_query, 1)
    expected = [{
        'Price': [24.49, 39.99],
        'Title': "BinGoDug Men's Basketball Shoes, Men's Fashion Sneakers, Air "
        'Basketball Shoes for Men, Womens Basketball Shoes, Mens Basketball '
        'Shoes, Boys Basketball Shoes, Youth Basketball Shoes Men Women',
        'asin': 'B09GKFNQWT'
    }, {
        'Price': [1.89, 7.58],
        'Title': "RQWEIN Comfortable Mesh Sneakers Men's Roading Running Shoes "
        'Tennis Shoes Casual Fashion Sneakers Outdoor Non Slip Gym Athletic '
        'Sport Shoes',
        'asin': 'B09BFY2R3R'
    }, {
        'Price': 100.0,
        'Title': 'PMUYBHF Womens Fashion Flat Shoes Comfortable Running Shoes '
        'Sneakers Tennis Athletic Shoe Casual Walking Shoes',
        'asin': 'B09P87V3LZ'
    }, {
        'Price': 100.0,
        'Title': 'PMUYBHF Fashion Travel Shoes Jogging Walking Sneakers Air Cushion '
        'Platform Loafers Air Cushion Mesh Shoes Walking Dance Shoes',
        'asin': 'B09N6SNKC1'
    }, {
        'Price': 100.0,
        'Title': "PMUYBHF Women's Ballet Flats Walking Flats Shoes Dressy Work Low "
        'Wedge Arch Suport Flats Shoes Slip On Dress Shoes',
        'asin': 'B09N6X5S74'
    }, {
        'Price': 100.0,
        'Title': "PWKSELW High-top Men's Basketball Shoes Outdoor Sports Shoes "
        'Cushioning Training Shoes Casual Running Shoes',
        'asin': 'B09MDB9V5W'
    }, {
        'Price': 100.0,
        'Title': "Women's Flat Shoes Classic Round Toe Slip Office Black Ballet "
        'Flats Walking Flats Shoes Casual Ballet Flats',
        'asin': 'B09N6PDFRF'
    }, {
        'Price': 100.0,
        'Title': "Women's Mid-Calf Boots Wide Calf Boots for Women Fashion Zipper "
        'Womens Shoes Pu Leather Casual Boots Womens Slip-On Womens Flat '
        "Shoes Med Heel Womens' Boots Winter Snow Boot Comfy Boots(,5.5)",
        'asin': 'B09N8ZHFNM'
    }, {
        'Price': 100.0,
        'Title': 'PMUYBHF Womens Leisure Fitness Running Sport Warm Sneakers Shoes '
        'Slip-On Mule Sneakers Womens Mules',
        'asin': 'B09P87DWGR'
    }, {
        'Price': 100.0,
        'Title': 'Men Dress Shoes Leather Modern Classic Business Shoes Lace Up '
        'Classic Office Shoes Business Formal Shoes for Men',
        'asin': 'B09R9MMTKR'
    }]
    assert output == expected

def test_convert_dict_to_actions():
    # Test RESULTS page type
    asin = "334490012932"
    page_num = 2
    products = [{
        'asin': '125331076844',
        'Title': 'Modern Tall Torchiere Floor Lamp Brushed Nickel Chrome Metal Decor Living Room',
        'Price': '$129.95'
    }, {
        'asin': '125109985453',
        'Title': 'Floor Lamps Set of 2 Polished Steel Crystal Glass for Living Room Bedroom',
        'Price': '$179.99'
    }, {
        'asin': '125265434055',
        'Title': 'Floor Lamp Nickel/Polished Concrete Finish with Off-White Linen Fabric Shade',
        'Price': '$130.68'
    }, {
        'asin': '195197281169',
        'Title': 'New ListingVintage Mid Century Modern Glass Amber Globe Tension Pole Floor Lamp Light',


        'asin': 'B09N6PDFRF'
    }, {
        'Price': 100.0,
        'Title': "Women's Mid-Calf Boots Wide Calf Boots for Women Fashion Zipper "
        'Womens Shoes Pu Leather Casual Boots Womens Slip-On Womens Flat '
        "Shoes Med Heel Womens' Boots Winter Snow Boot Comfy Boots(,5.5)",
        'asin': 'B09N8ZHFNM'
    }, {
        'Price': 100.0,
        'Title': 'PMUYBHF Womens Leisure Fitness Running Sport Warm Sneakers Shoes '
        'Slip-On Mule Sneakers Womens Mules',
        'asin': 'B09P87DWGR'
    }, {
        'Price': 100.0,
        'Title': 'Men Dress Shoes Leather Modern Classic Business Shoes Lace Up '
        'Classic Office Shoes Business Formal Shoes for Men',
        'asin': 'B09R9MMTKR'
    }]
    assert output == expected

def test_convert_dict_to_actions():
    # Test RESULTS page type
    asin = "334490012932"
    page_num = 2
    products = [{
        'asin': '125331076844',
        'Title': 'Modern Tall Torchiere Floor Lamp Brushed Nickel Chrome Metal Decor Living Room',
        'Price': '$129.95'
    }, {
        'asin': '125109985453',
        'Title': 'Floor Lamps Set of 2 Polished Steel Crystal Glass for Living Room Bedroom',
        'Price': '$179.99'
    }, {
        'asin': '125265434055',
        'Title': 'Floor Lamp Nickel/Polished Concrete Finish with Off-White Linen Fabric Shade',
        'Price': '$130.68'
    }, {
        'asin': '195197281169',
        'Title': 'New ListingVintage Mid Century Modern Glass Amber Globe Tension Pole Floor Lamp Light',
        'Price': '$165.00'
    }, {
        'asin': '195197512929',
        'Title': 'New ListingVTG Brass Floor Lamp Glass Shade 63.5" Tall 12" Diameter Glass Shade Original',
        'Price': '$279.45'
    }, {
        'asin': '304550250934',
        'Title': 'Vintage Mid Century Modern 3 Light Tension Pole Floor Lamp glass shades atomic a',
        'Price': '$149.99'
    }, {
        'asin': '175338033811',
        'Title': 'Antique FOSTORIA Ornate Brass Piano  Adjustable Floor Oil Lamp up to 76" Tall !!',
        'Price': '$1,995.00'
    }, {
        'asin': '334490012932',
        'Title': 'Vintage Mid Century Glass Shade Amber Globe 3 Tension Pole Floor Lamp Light MCM',
        'Price': '$128.00'
    }, {
        'asin': '185433933521',
        'Title': 'Brass & Pink Glass Lotus 6 Petal Lamp Shades Set Of Two Replacement Parts As Is',
        'Price': '$90.00'
    }]

    actions = convert_dict_to_actions(Page.RESULTS, products, asin, page_num)

    assert actions['valid'] == [
        'click[back to search]',
        'click[< prev]',
        'click[item - Modern Tall Torchiere Floor Lamp Brushed Nickel Chrome Metal Decor Living Room]',
        'click[item - Floor Lamps Set of 2 Polished Steel Crystal Glass for Living Room Bedroom]',
        'click[item - Floor Lamp Nickel/Polished Concrete Finish with Off-White Linen Fabric Shade]',
        'click[item - New ListingVintage Mid Century Modern Glass Amber Globe Tension Pole Floor Lamp Light]',
        'click[item - New ListingVTG Brass Floor Lamp Glass Shade 63.5" Tall 12" Diameter Glass Shade Original]',
        'click[item - Vintage Mid Century Modern 3 Light Tension Pole Floor Lamp glass shades atomic a]',
        'click[item - Antique FOSTORIA Ornate Brass Piano  Adjustable Floor Oil Lamp up to 76" Tall !!]',
        'click[item - Vintage Mid Century Glass Shade Amber Globe 3 Tension Pole Floor Lamp Light MCM]',
        'click[item - Brass & Pink Glass Lotus 6 Petal Lamp Shades Set Of Two Replacement Parts As Is]'
    ]

    # Test ITEM_PAGE page type
    asin = "224636269803"
    products = {
        '224636269803': {
            'asin': '224636269803',
            'Title': 'Sony SRS-XB01 EXTRA BASS Portable Water-Resistant  Wireless Bluetooth Speaker',
            'Price': '24.99',
            'MainImage': 'https://i.ebayimg.com/images/g/jVEAAOSwCLBhXLuD/s-l500.jpg',
            'Rating': None,
            'options': {
                'Color': ['Black', 'White', 'Red', 'Blue']
            },
            'option_to_image': {},
            'Description': "eBay Sony EXTRA BASS Portable Water-Resistant Wireless Bluetooth SpeakerBRAND NEW ITEMFREE SHIPPING WITHIN USA30 DAY RETURN POLICYKey FeaturesEXTRA BASS for deep, punchy soundCompact portable designUp to 6 hours of battery lifeWater resistant for worry-free useSupplied with color-coordinated strap What's in the Box?Sony EXTRA BASS Portable Bluetooth SpeakerPower supplyUser manual HIGHLIGHTSMUSIC THAT TRAVELSSmall size but mighty in volume to deliver powerful beats wherever you travelHANDS FREE CALLINGWith the built-in microphone, taking calls from your smartphone is easy. SPLASHPROOF CASINGTake to the pool or beach without worrying about water damaging the speaker unit UPGRADE THE AUDIOWirelessly connects 2 speakers and achieve stereo sound with speaker add function LONGER BATTERY LIFELonger Virtual Happy Hours with this rechargeable speaker's 6 hour battery life Technical SpecsFeatureValueBrandSonyTypePortable speakerModel NumberSRSXB01BluetoothYesFrequency range2.4 GHzMax. Communication Range32 ftBattery LifeApprox. 6 hrsWater ProtectionIPX5Input and Output TerminalsStereo Mini Jack (IN)Dimensions (W x H x D)Approx. 3 1/4 × 2 3/8 × 2 1/4 inWeightApprox. 5.65 oz",
            'BulletPoints': "Item specifics Condition:New: A brand-new, unused, unopened, undamaged item in its original packaging (where packaging is ...  Read moreabout the conditionNew: A brand-new, unused, unopened, undamaged item in its original packaging (where packaging is applicable). Packaging should be the same as what is found in a retail store, unless the item is handmade or was packaged by the manufacturer in non-retail packaging, such as an unprinted box or plastic bag. See the seller's listing for full details. See all condition definitionsopens in a new window or tab  Model:EXTRA BASS Connectivity:Bluetooth, Wireless Type:Portable Speaker System Compatible Model:EXTRA BASS, Portable Water-Resistant Features:Bluetooth, Water-Resistant MPN:SRS-XB01/B, SRS-XB01/L, SRS-XB01/R, SRS-XB01/W Brand:Sony"
        }
    }

    actions = convert_dict_to_actions(Page.ITEM_PAGE, products, asin, 1)

    assert actions['valid'] == ['click[back to search]', 'click[< prev]', 'click[description]', 'click[features]', 'click[buy now]', 'click[Black]', 'click[White]', 'click[Red]', 'click[Blue]']

    # Test SUB_PAGE page type
    actions = convert_dict_to_actions(Page.SUB_PAGE, {}, "12345", 1)

    assert actions['valid'] == ['click[back to search]', 'click[< prev]']


"""
Test the text gym environment.

TODO: move to testing dir for more rigorous tests
"""
import gym
from rich import print
from rich.markup import escape

from web_agent_site.envs import WebAgentTextEnv
from web_agent_site.models import RandomPolicy
from web_agent_site.utils import DEBUG_PROD_SIZE

if __name__ == '__main__':
    env = gym.make('WebAgentTextEnv-v0', observation_mode='text', num_products=DEBUG_PROD_SIZE)
    env.reset()

    try:
        policy = RandomPolicy()

        observation = env.observation
        while True:
            print(observation)
            available_actions = env.get_available_actions()
            print('Available actions:', available_actions)
            # action = policy.forward(observation, available_actions)
            action = "search[ flat screen tv under $500 ]"
            observation, reward, done, info = env.step(action)
            print("observation:", observation)
            print("info:", info)
            print(f'Taking action "{escape(action)}" -> Reward = {reward}')
            if done:
                break
    finally:
        env.close()


"""
Test the site gym environment.

TODO: move to testing dir for more rigorous tests
"""
import gym
from rich import print
from rich.markup import escape

from web_agent_site.envs import WebAgentSiteEnv
from web_agent_site.models import (
    HumanPolicy,
    RandomPolicy,
)
from web_agent_site.utils import DEBUG_PROD_SIZE

if __name__ == '__main__':
    #env = gym.make('WebAgentSite-v0')
    #env = WebAgentSiteEnv(render=True, pause=2.0)
    #env = WebAgentSiteEnv(observation_mode='html', render=False)
    env = WebAgentSiteEnv(observation_mode='text', render=False, num_products=DEBUG_PROD_SIZE)
    global_step = 0

    try:
        #policy = HumanPolicy()
        policy = RandomPolicy()

        observation = env.observation
        while True:
            print(observation)
            available_actions = env.get_available_actions()
            print('Available actions:', available_actions)
            action = policy.forward(observation, available_actions)
            observation, reward, done, info = env.step(action)
            print(f'Taking action "{escape(action)}" -> Reward = {reward}')
            if done:
                break
            global_step += 1
    finally:
        env.close()


import argparse, json, logging, random
from pathlib import Path
from ast import literal_eval

from flask import (
    Flask,
    request,
    redirect,
    url_for
)

from rich import print

from web_agent_site.engine.engine import (
    load_products,
    init_search_engine,
    convert_web_app_string_to_var,
    get_top_n_product_from_keywords,
    get_product_per_page,
    map_action_to_html,
    END_BUTTON
)
from web_agent_site.engine.goal import get_reward, get_goals
from web_agent_site.utils import (
    generate_mturk_code,
    setup_logger,
    DEFAULT_FILE_PATH,
    DEBUG_PROD_SIZE,
)

app = Flask(__name__)

search_engine = None
all_products = None
product_item_dict = None
product_prices = None
attribute_to_asins = None
goals = None
weights = None

user_sessions = dict()
user_log_dir = None
SHOW_ATTRS_TAB = False

@app.route('/')
def home():
    return redirect(url_for('index', session_id="abc"))

@app.route('/<session_id>', methods=['GET', 'POST'])
def index(session_id):
    global user_log_dir
    global all_products, product_item_dict, \
           product_prices, attribute_to_asins, \
           search_engine, \
           goals, weights, user_sessions

    if search_engine is None:
        all_products, product_item_dict, product_prices, attribute_to_asins = \
            load_products(
                filepath=DEFAULT_FILE_PATH,
                num_products=DEBUG_PROD_SIZE
            )
        search_engine = init_search_engine(num_products=DEBUG_PROD_SIZE)
        goals = get_goals(all_products, product_prices)
        # random.seed(233)
        # random.shuffle(goals)
        weights = [goal['weight'] for goal in goals]

    if session_id not in user_sessions and 'fixed' in session_id:
        goal_dix = int(session_id.split('_')[-1])
        goal = goals[goal_dix]
        instruction_text = goal['instruction_text']
        user_sessions[session_id] = {'goal': goal, 'done': False}
        if user_log_dir is not None:
            setup_logger(session_id, user_log_dir)
    elif session_id not in user_sessions:
        goal = random.choices(goals, weights)[0]
        instruction_text = goal['instruction_text']
        user_sessions[session_id] = {'goal': goal, 'done': False}
        if user_log_dir is not None:
            setup_logger(session_id, user_log_dir)
    else:
        instruction_text = user_sessions[session_id]['goal']['instruction_text']

    if request.method == 'POST' and 'search_query' in request.form:
        keywords = request.form['search_query'].lower().split(' ')
        return redirect(url_for(
            'search_results',
            session_id=session_id,
            keywords=keywords,
            page=1,
        ))
    if user_log_dir is not None:
        logger = logging.getLogger(session_id)
        logger.info(json.dumps(dict(
            page='index',
            url=request.url,
            goal=user_sessions[session_id]['goal'],
        )))
    return map_action_to_html(
        'start',
        session_id=session_id,
        instruction_text=instruction_text,
    )

@app.route(
    '/search_results/<session_id>/<keywords>/<page>',
    methods=['GET', 'POST']
)
def search_results(session_id, keywords, page):
    instruction_text = user_sessions[session_id]['goal']['instruction_text']
    page = convert_web_app_string_to_var('page', page)
    keywords = convert_web_app_string_to_var('keywords', keywords)
    top_n_products = get_top_n_product_from_keywords(
        keywords,
        search_engine,
        all_products,
        product_item_dict,
        attribute_to_asins,
    )
    products = get_product_per_page(top_n_products, page)
    html = map_action_to_html(
        'search',
        session_id=session_id,
        products=products,
        keywords=keywords,
        page=page,
        total=len(top_n_products),
        instruction_text=instruction_text,
    )
    logger = logging.getLogger(session_id)
    logger.info(json.dumps(dict(
        page='search_results',
        url=request.url,
        goal=user_sessions[session_id]['goal'],
        content=dict(
            keywords=keywords,
            search_result_asins=[p['asin'] for p in products],
            page=page,
        )
    )))
    return html

@app.route(
    '/item_page/<session_id>/<asin>/<keywords>/<page>/<options>',
    methods=['GET', 'POST']
)
def item_page(session_id, asin, keywords, page, options):
    options = literal_eval(options)
    product_info = product_item_dict[asin]

    goal_instruction = user_sessions[session_id]['goal']['instruction_text']
    product_info['goal_instruction'] = goal_instruction

    html = map_action_to_html(
        'click',
        session_id=session_id,
        product_info=product_info,
        keywords=keywords,
        page=page,
        asin=asin,
        options=options,
        instruction_text=goal_instruction,
        show_attrs=SHOW_ATTRS_TAB,
    )
    logger = logging.getLogger(session_id)
    logger.info(json.dumps(dict(
        page='item_page',
        url=request.url,
        goal=user_sessions[session_id]['goal'],
        content=dict(
            keywords=keywords,
            page=page,
            asin=asin,
            options=options,
        )
    )))
    return html

@app.route(
    '/item_sub_page/<session_id>/<asin>/<keywords>/<page>/<sub_page>/<options>',
    methods=['GET', 'POST']
)
def item_sub_page(session_id, asin, keywords, page, sub_page, options):
    options = literal_eval(options)
    product_info = product_item_dict[asin]

    goal_instruction = user_sessions[session_id]['goal']['instruction_text']
    product_info['goal_instruction'] = goal_instruction

    html = map_action_to_html(
        f'click[{sub_page}]',
        session_id=session_id,
        product_info=product_info,
        keywords=keywords,
        page=page,
        asin=asin,
        options=options,
        instruction_text=goal_instruction
    )
    logger = logging.getLogger(session_id)
    logger.info(json.dumps(dict(
        page='item_sub_page',
        url=request.url,
        goal=user_sessions[session_id]['goal'],
        content=dict(
            keywords=keywords,
            page=page,
            asin=asin,
            options=options,
        )
    )))
    return html

@app.route('/done/<session_id>/<asin>/<options>', methods=['GET', 'POST'])
def done(session_id, asin, options):
    options = literal_eval(options)
    goal = user_sessions[session_id]['goal']
    purchased_product = product_item_dict[asin]
    price = product_prices[asin]

    reward, reward_info = get_reward(
        purchased_product,
        goal,
        price=price,
        options=options,
        verbose=True
    )
    user_sessions[session_id]['done'] = True
    user_sessions[session_id]['reward'] = reward
    print(user_sessions)

    logger = logging.getLogger(session_id)
    logger.info(json.dumps(dict(
        page='done',
        url=request.url,
        goal=goal,
        content=dict(
            asin=asin,
            options=options,
            price=price,
        ),
        reward=reward,
        reward_info=reward_info,
    )))
    del logging.root.manager.loggerDict[session_id]

    return map_action_to_html(
        f'click[{END_BUTTON}]',
        session_id=session_id,
        reward=reward,
        asin=asin,
        options=options,
        reward_info=reward_info,
        query=purchased_product['query'],
        category=purchased_product['category'],
        product_category=purchased_product['product_category'],
        goal_attrs=user_sessions[session_id]['goal']['attributes'],
        purchased_attrs=purchased_product['Attributes'],
        goal=goal,
        mturk_code=generate_mturk_code(session_id),
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="WebShop flask app backend configuration")
    parser.add_argument("--log", action='store_true', help="Log actions on WebShop in trajectory file")
    parser.add_argument("--attrs", action='store_true', help="Show attributes tab in item page")

    args = parser.parse_args()
    if args.log:
        user_log_dir = Path('user_session_logs/mturk')
        user_log_dir.mkdir(parents=True, exist_ok=True)
    SHOW_ATTRS_TAB = args.attrs

    app.run(host='0.0.0.0', port=3000)


home

    setup_logger,
    DEFAULT_FILE_PATH,
    DEBUG_PROD_SIZE,
)

app = Flask(__name__)

search_engine = None
all_products = None
product_item_dict = None
product_prices = None
attribute_to_asins = None
goals = None
weights = None

user_sessions = dict()
user_log_dir = None
SHOW_ATTRS_TAB = False

@app.route('/')
def home():
    return redirect(url_for('index', session_id="abc"))

@app.route('/<session_id>', methods=['GET', 'POST'])
def index(session_id):
    global user_log_dir
    global all_products, product_item_dict, \
           product_prices, attribute_to_asins, \
           search_engine, \
           goals, weights, user_sessions

    if search_engine is None:
        all_products, product_item_dict, product_prices, attribute_to_asins = \
            load_products(
                filepath=DEFAULT_FILE_PATH,
                num_products=DEBUG_PROD_SIZE
            )
        search_engine = init_search_engine(num_products=DEBUG_PROD_SIZE)
        goals = get_goals(all_products, product_prices)
        # random.seed(233)
        # random.shuffle(goals)
        weights = [goal['weight'] for goal in goals]


index

app = Flask(__name__)

search_engine = None
all_products = None
product_item_dict = None
product_prices = None
attribute_to_asins = None
goals = None
weights = None

user_sessions = dict()
user_log_dir = None
SHOW_ATTRS_TAB = False

@app.route('/')
def home():
    return redirect(url_for('index', session_id="abc"))

@app.route('/<session_id>', methods=['GET', 'POST'])
def index(session_id):
    global user_log_dir
    global all_products, product_item_dict, \
           product_prices, attribute_to_asins, \
           search_engine, \
           goals, weights, user_sessions

    if search_engine is None:
        all_products, product_item_dict, product_prices, attribute_to_asins = \
            load_products(
                filepath=DEFAULT_FILE_PATH,
                num_products=DEBUG_PROD_SIZE
            )
        search_engine = init_search_engine(num_products=DEBUG_PROD_SIZE)
        goals = get_goals(all_products, product_prices)
        # random.seed(233)
        # random.shuffle(goals)
        weights = [goal['weight'] for goal in goals]

    if session_id not in user_sessions and 'fixed' in session_id:
        goal_dix = int(session_id.split('_')[-1])
        goal = goals[goal_dix]
        instruction_text = goal['instruction_text']
        user_sessions[session_id] = {'goal': goal, 'done': False}
        if user_log_dir is not None:
            setup_logger(session_id, user_log_dir)
    elif session_id not in user_sessions:
        goal = random.choices(goals, weights)[0]
        instruction_text = goal['instruction_text']
        user_sessions[session_id] = {'goal': goal, 'done': False}
        if user_log_dir is not None:
            setup_logger(session_id, user_log_dir)
    else:
        instruction_text = user_sessions[session_id]['goal']['instruction_text']

    if request.method == 'POST' and 'search_query' in request.form:
        keywords = request.form['search_query'].lower().split(' ')
        return redirect(url_for(
            'search_results',
            session_id=session_id,
            keywords=keywords,
            page=1,
        ))
    if user_log_dir is not None:
        logger = logging.getLogger(session_id)
        logger.info(json.dumps(dict(
            page='index',
            url=request.url,
            goal=user_sessions[session_id]['goal'],
        )))
    return map_action_to_html(
        'start',
        session_id=session_id,
        instruction_text=instruction_text,
    )

@app.route(
    '/search_results/<session_id>/<keywords>/<page>',
    methods=['GET', 'POST']
)
def search_results(session_id, keywords, page):
    instruction_text = user_sessions[session_id]['goal']['instruction_text']
    page = convert_web_app_string_to_var('page', page)
    keywords = convert_web_app_string_to_var('keywords', keywords)
    top_n_products = get_top_n_product_from_keywords(
        keywords,
        search_engine,
        all_products,
        product_item_dict,
        attribute_to_asins,
    )
    products = get_product_per_page(top_n_products, page)
    html = map_action_to_html(
        'search',


            page=1,
        ))
    if user_log_dir is not None:
        logger = logging.getLogger(session_id)
        logger.info(json.dumps(dict(
            page='index',
            url=request.url,
            goal=user_sessions[session_id]['goal'],
        )))
    return map_action_to_html(
        'start',
        session_id=session_id,
        instruction_text=instruction_text,
    )

@app.route(
    '/search_results/<session_id>/<keywords>/<page>',
    methods=['GET', 'POST']
)
def search_results(session_id, keywords, page):
    instruction_text = user_sessions[session_id]['goal']['instruction_text']
    page = convert_web_app_string_to_var('page', page)
    keywords = convert_web_app_string_to_var('keywords', keywords)
    top_n_products = get_top_n_product_from_keywords(
        keywords,
        search_engine,
        all_products,
        product_item_dict,
        attribute_to_asins,
    )
    products = get_product_per_page(top_n_products, page)
    html = map_action_to_html(
        'search',
        session_id=session_id,
        products=products,
        keywords=keywords,
        page=page,
        total=len(top_n_products),
        instruction_text=instruction_text,
    )
    logger = logging.getLogger(session_id)
    logger.info(json.dumps(dict(
        page='search_results',
        url=request.url,
        goal=user_sessions[session_id]['goal'],
        content=dict(
            keywords=keywords,
            search_result_asins=[p['asin'] for p in products],
            page=page,
        )
    )))
    return html

@app.route(
    '/item_page/<session_id>/<asin>/<keywords>/<page>/<options>',
    methods=['GET', 'POST']
)
def item_page(session_id, asin, keywords, page, options):
    options = literal_eval(options)
    product_info = product_item_dict[asin]

    goal_instruction = user_sessions[session_id]['goal']['instruction_text']
    product_info['goal_instruction'] = goal_instruction

    html = map_action_to_html(
        'click',
        session_id=session_id,
        product_info=product_info,
        keywords=keywords,
        page=page,
        asin=asin,


        instruction_text=instruction_text,
    )
    logger = logging.getLogger(session_id)
    logger.info(json.dumps(dict(
        page='search_results',
        url=request.url,
        goal=user_sessions[session_id]['goal'],
        content=dict(
            keywords=keywords,
            search_result_asins=[p['asin'] for p in products],
            page=page,
        )
    )))
    return html

@app.route(
    '/item_page/<session_id>/<asin>/<keywords>/<page>/<options>',
    methods=['GET', 'POST']
)
def item_page(session_id, asin, keywords, page, options):
    options = literal_eval(options)
    product_info = product_item_dict[asin]

    goal_instruction = user_sessions[session_id]['goal']['instruction_text']
    product_info['goal_instruction'] = goal_instruction

    html = map_action_to_html(
        'click',
        session_id=session_id,
        product_info=product_info,
        keywords=keywords,
        page=page,
        asin=asin,
        options=options,
        instruction_text=goal_instruction,
        show_attrs=SHOW_ATTRS_TAB,
    )
    logger = logging.getLogger(session_id)
    logger.info(json.dumps(dict(
        page='item_page',
        url=request.url,
        goal=user_sessions[session_id]['goal'],
        content=dict(
            keywords=keywords,
            page=page,
            asin=asin,
            options=options,
        )
    )))
    return html

@app.route(
    '/item_sub_page/<session_id>/<asin>/<keywords>/<page>/<sub_page>/<options>',
    methods=['GET', 'POST']
)
def item_sub_page(session_id, asin, keywords, page, sub_page, options):
    options = literal_eval(options)
    product_info = product_item_dict[asin]

    goal_instruction = user_sessions[session_id]['goal']['instruction_text']
    product_info['goal_instruction'] = goal_instruction

    html = map_action_to_html(
        f'click[{sub_page}]',
        session_id=session_id,
        product_info=product_info,
        keywords=keywords,
        page=page,
        asin=asin,


    )
    logger = logging.getLogger(session_id)
    logger.info(json.dumps(dict(
        page='item_page',
        url=request.url,
        goal=user_sessions[session_id]['goal'],
        content=dict(
            keywords=keywords,
            page=page,
            asin=asin,
            options=options,
        )
    )))
    return html

@app.route(
    '/item_sub_page/<session_id>/<asin>/<keywords>/<page>/<sub_page>/<options>',
    methods=['GET', 'POST']
)
def item_sub_page(session_id, asin, keywords, page, sub_page, options):
    options = literal_eval(options)
    product_info = product_item_dict[asin]

    goal_instruction = user_sessions[session_id]['goal']['instruction_text']
    product_info['goal_instruction'] = goal_instruction

    html = map_action_to_html(
        f'click[{sub_page}]',
        session_id=session_id,
        product_info=product_info,
        keywords=keywords,
        page=page,
        asin=asin,
        options=options,
        instruction_text=goal_instruction
    )
    logger = logging.getLogger(session_id)
    logger.info(json.dumps(dict(
        page='item_sub_page',
        url=request.url,
        goal=user_sessions[session_id]['goal'],
        content=dict(
            keywords=keywords,
            page=page,
            asin=asin,
            options=options,
        )
    )))
    return html

@app.route('/done/<session_id>/<asin>/<options>', methods=['GET', 'POST'])
def done(session_id, asin, options):
    options = literal_eval(options)
    goal = user_sessions[session_id]['goal']
    purchased_product = product_item_dict[asin]
    price = product_prices[asin]

    reward, reward_info = get_reward(
        purchased_product,
        goal,
        price=price,
        options=options,
        verbose=True
    )
    user_sessions[session_id]['done'] = True
    user_sessions[session_id]['reward'] = reward
    print(user_sessions)


done

        asin=asin,
        options=options,
        instruction_text=goal_instruction
    )
    logger = logging.getLogger(session_id)
    logger.info(json.dumps(dict(
        page='item_sub_page',
        url=request.url,
        goal=user_sessions[session_id]['goal'],
        content=dict(
            keywords=keywords,
            page=page,
            asin=asin,
            options=options,
        )
    )))
    return html

@app.route('/done/<session_id>/<asin>/<options>', methods=['GET', 'POST'])
def done(session_id, asin, options):
    options = literal_eval(options)
    goal = user_sessions[session_id]['goal']
    purchased_product = product_item_dict[asin]
    price = product_prices[asin]

    reward, reward_info = get_reward(
        purchased_product,
        goal,
        price=price,
        options=options,
        verbose=True
    )
    user_sessions[session_id]['done'] = True
    user_sessions[session_id]['reward'] = reward
    print(user_sessions)

    logger = logging.getLogger(session_id)
    logger.info(json.dumps(dict(
        page='done',
        url=request.url,
        goal=goal,
        content=dict(
            asin=asin,
            options=options,
            price=price,
        ),
        reward=reward,
        reward_info=reward_info,
    )))
    del logging.root.manager.loggerDict[session_id]

    return map_action_to_html(
        f'click[{END_BUTTON}]',
        session_id=session_id,
        reward=reward,
        asin=asin,
        options=options,
        reward_info=reward_info,
        query=purchased_product['query'],
        category=purchased_product['category'],
        product_category=purchased_product['product_category'],
        goal_attrs=user_sessions[session_id]['goal']['attributes'],
        purchased_attrs=purchased_product['Attributes'],
        goal=goal,
        mturk_code=generate_mturk_code(session_id),
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="WebShop flask app backend configuration")
    parser.add_argument("--log", action='store_true', help="Log actions on WebShop in trajectory file")
    parser.add_argument("--attrs", action='store_true', help="Show attributes tab in item page")

    args = parser.parse_args()
    if args.log:
        user_log_dir = Path('user_session_logs/mturk')
        user_log_dir.mkdir(parents=True, exist_ok=True)
    SHOW_ATTRS_TAB = args.attrs

    app.run(host='0.0.0.0', port=3000)


import bisect
import hashlib
import logging
import random
from os.path import dirname, abspath, join

BASE_DIR = dirname(abspath(__file__))
DEBUG_PROD_SIZE = None  # set to `None` to disable

# DEFAULT_ATTR_PATH = join(BASE_DIR, '../data/items_ins_v2_1000.json')
# DEFAULT_FILE_PATH = join(BASE_DIR, '../data/items_shuffle_1000.json')
DEFAULT_ATTR_PATH = join(BASE_DIR, '../data/items_ins_v2.json')
DEFAULT_FILE_PATH = join(BASE_DIR, '../data/items_shuffle.json')
DEFAULT_REVIEW_PATH = join(BASE_DIR, '../data/reviews.json')

FEAT_CONV = join(BASE_DIR, '../data/feat_conv.pt')
FEAT_IDS = join(BASE_DIR, '../data/feat_ids.pt')

HUMAN_ATTR_PATH = join(BASE_DIR, '../data/new_human_ins_sampled_1.json')

# HUMAN_ATTR_PATH = join(BASE_DIR, '../data/items_human_ins.json')

def random_idx(cum_weights):
    """Generate random index by sampling uniformly from sum of all weights, then
    selecting the `min` between the position to keep the list sorted (via bisect)
    and the value of the second to last index
    """
    pos = random.uniform(0, cum_weights[-1])
    idx = bisect.bisect(cum_weights, pos)
    idx = min(idx, len(cum_weights) - 2)
    return idx

def setup_logger(session_id, user_log_dir):
    """Creates a log file and logging object for the corresponding session ID"""
    logger = logging.getLogger(session_id)
    formatter = logging.Formatter('%(message)s')
    file_handler = logging.FileHandler(
        user_log_dir / f'{session_id}.jsonl',
        mode='w'
    )
    file_handler.setFormatter(formatter)
    logger.setLevel(logging.INFO)
    logger.addHandler(file_handler)
    return logger

def generate_mturk_code(session_id: str) -> str:
    """Generates a redeem code corresponding to the session ID for an MTurk
    worker once the session is completed
    """
    sha = hashlib.sha1(session_id.encode())
    return sha.hexdigest()[:10].upper()


import logging
import random
from os.path import dirname, abspath, join

BASE_DIR = dirname(abspath(__file__))
DEBUG_PROD_SIZE = None  # set to `None` to disable

# DEFAULT_ATTR_PATH = join(BASE_DIR, '../data/items_ins_v2_1000.json')
# DEFAULT_FILE_PATH = join(BASE_DIR, '../data/items_shuffle_1000.json')
DEFAULT_ATTR_PATH = join(BASE_DIR, '../data/items_ins_v2.json')
DEFAULT_FILE_PATH = join(BASE_DIR, '../data/items_shuffle.json')
DEFAULT_REVIEW_PATH = join(BASE_DIR, '../data/reviews.json')

FEAT_CONV = join(BASE_DIR, '../data/feat_conv.pt')
FEAT_IDS = join(BASE_DIR, '../data/feat_ids.pt')

HUMAN_ATTR_PATH = join(BASE_DIR, '../data/new_human_ins_sampled_1.json')

# HUMAN_ATTR_PATH = join(BASE_DIR, '../data/items_human_ins.json')

def random_idx(cum_weights):
    """Generate random index by sampling uniformly from sum of all weights, then
    selecting the `min` between the position to keep the list sorted (via bisect)
    and the value of the second to last index
    """
    pos = random.uniform(0, cum_weights[-1])
    idx = bisect.bisect(cum_weights, pos)
    idx = min(idx, len(cum_weights) - 2)
    return idx

def setup_logger(session_id, user_log_dir):
    """Creates a log file and logging object for the corresponding session ID"""
    logger = logging.getLogger(session_id)
    formatter = logging.Formatter('%(message)s')
    file_handler = logging.FileHandler(
        user_log_dir / f'{session_id}.jsonl',
        mode='w'
    )
    file_handler.setFormatter(formatter)
    logger.setLevel(logging.INFO)
    logger.addHandler(file_handler)
    return logger

def generate_mturk_code(session_id: str) -> str:
    """Generates a redeem code corresponding to the session ID for an MTurk
    worker once the session is completed
    """
    sha = hashlib.sha1(session_id.encode())
    return sha.hexdigest()[:10].upper()


DEFAULT_FILE_PATH = join(BASE_DIR, '../data/items_shuffle.json')
DEFAULT_REVIEW_PATH = join(BASE_DIR, '../data/reviews.json')

FEAT_CONV = join(BASE_DIR, '../data/feat_conv.pt')
FEAT_IDS = join(BASE_DIR, '../data/feat_ids.pt')

HUMAN_ATTR_PATH = join(BASE_DIR, '../data/new_human_ins_sampled_1.json')

# HUMAN_ATTR_PATH = join(BASE_DIR, '../data/items_human_ins.json')

def random_idx(cum_weights):
    """Generate random index by sampling uniformly from sum of all weights, then
    selecting the `min` between the position to keep the list sorted (via bisect)
    and the value of the second to last index
    """
    pos = random.uniform(0, cum_weights[-1])
    idx = bisect.bisect(cum_weights, pos)
    idx = min(idx, len(cum_weights) - 2)
    return idx

def setup_logger(session_id, user_log_dir):
    """Creates a log file and logging object for the corresponding session ID"""
    logger = logging.getLogger(session_id)
    formatter = logging.Formatter('%(message)s')
    file_handler = logging.FileHandler(
        user_log_dir / f'{session_id}.jsonl',
        mode='w'
    )
    file_handler.setFormatter(formatter)
    logger.setLevel(logging.INFO)
    logger.addHandler(file_handler)
    return logger

def generate_mturk_code(session_id: str) -> str:
    """Generates a redeem code corresponding to the session ID for an MTurk
    worker once the session is completed
    """
    sha = hashlib.sha1(session_id.encode())
    return sha.hexdigest()[:10].upper()


    and the value of the second to last index
    """
    pos = random.uniform(0, cum_weights[-1])
    idx = bisect.bisect(cum_weights, pos)
    idx = min(idx, len(cum_weights) - 2)
    return idx

def setup_logger(session_id, user_log_dir):
    """Creates a log file and logging object for the corresponding session ID"""
    logger = logging.getLogger(session_id)
    formatter = logging.Formatter('%(message)s')
    file_handler = logging.FileHandler(
        user_log_dir / f'{session_id}.jsonl',
        mode='w'
    )
    file_handler.setFormatter(formatter)
    logger.setLevel(logging.INFO)
    logger.addHandler(file_handler)
    return logger

def generate_mturk_code(session_id: str) -> str:
    """Generates a redeem code corresponding to the session ID for an MTurk
    worker once the session is completed
    """
    sha = hashlib.sha1(session_id.encode())
    return sha.hexdigest()[:10].upper()


"""
Model implementations. The model interface should be suitable for both
the ``site env'' and the ``text env''.
"""
import random

random.seed(4)

class BasePolicy:
    def __init__(self):
        pass

    def forward(observation, available_actions):
        """
        Args:
            observation (`str`):
                HTML string

            available_actions ():
                ...
        Returns:
            action (`str`):
                Return string of the format ``action_name[action_arg]''.
                Examples:
                    - search[white shoes]
                    - click[button=Reviews]
                    - click[button=Buy Now]
        """
        raise NotImplementedError

class HumanPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        action = input('> ')
        return action

class RandomPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        if available_actions['has_search_bar']:
            action = 'search[shoes]'
        else:
            action_arg = random.choice(available_actions['clickables'])
            action = f'click[{action_arg}]'
        return action


BasePolicy

"""
Model implementations. The model interface should be suitable for both
the ``site env'' and the ``text env''.
"""
import random

random.seed(4)

class BasePolicy:
    def __init__(self):
        pass

    def forward(observation, available_actions):
        """
        Args:
            observation (`str`):
                HTML string

            available_actions ():
                ...
        Returns:
            action (`str`):
                Return string of the format ``action_name[action_arg]''.
                Examples:
                    - search[white shoes]
                    - click[button=Reviews]
                    - click[button=Buy Now]
        """
        raise NotImplementedError

class HumanPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        action = input('> ')
        return action

class RandomPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        if available_actions['has_search_bar']:
            action = 'search[shoes]'
        else:
            action_arg = random.choice(available_actions['clickables'])


HumanPolicy

    def forward(observation, available_actions):
        """
        Args:
            observation (`str`):
                HTML string

            available_actions ():
                ...
        Returns:
            action (`str`):
                Return string of the format ``action_name[action_arg]''.
                Examples:
                    - search[white shoes]
                    - click[button=Reviews]
                    - click[button=Buy Now]
        """
        raise NotImplementedError

class HumanPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        action = input('> ')
        return action

class RandomPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        if available_actions['has_search_bar']:
            action = 'search[shoes]'
        else:
            action_arg = random.choice(available_actions['clickables'])
            action = f'click[{action_arg}]'
        return action


RandomPolicy

        Returns:
            action (`str`):
                Return string of the format ``action_name[action_arg]''.
                Examples:
                    - search[white shoes]
                    - click[button=Reviews]
                    - click[button=Buy Now]
        """
        raise NotImplementedError

class HumanPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        action = input('> ')
        return action

class RandomPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        if available_actions['has_search_bar']:
            action = 'search[shoes]'
        else:
            action_arg = random.choice(available_actions['clickables'])
            action = f'click[{action_arg}]'
        return action


            action (`str`):
                Return string of the format ``action_name[action_arg]''.
                Examples:
                    - search[white shoes]
                    - click[button=Reviews]
                    - click[button=Buy Now]
        """
        raise NotImplementedError

class HumanPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        action = input('> ')
        return action

class RandomPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        if available_actions['has_search_bar']:
            action = 'search[shoes]'
        else:
            action_arg = random.choice(available_actions['clickables'])
            action = f'click[{action_arg}]'
        return action


forward

                    - search[white shoes]
                    - click[button=Reviews]
                    - click[button=Buy Now]
        """
        raise NotImplementedError

class HumanPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        action = input('> ')
        return action

class RandomPolicy(BasePolicy):
    def __init__(self):
        super().__init__()

    def forward(self, observation, available_actions):
        if available_actions['has_search_bar']:
            action = 'search[shoes]'
        else:
            action_arg = random.choice(available_actions['clickables'])
            action = f'click[{action_arg}]'
        return action


import json
import yaml
import random
from pathlib import Path
from collections import defaultdict

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_extraction import text as sk_text
import pandas as pd
from tqdm import tqdm
from rich import print

ITEMS_PATH = './data/ITEMS_mar1.json'
REVIEWS_PATH = './data/reviews.json'
ATTR_DIR = './data/attributes'

random.seed(0)

def get_stop_words():
    extra_stop_words = set([str(i) for i in range(1000)])
    stop_words = sk_text.ENGLISH_STOP_WORDS.union(extra_stop_words)
    return stop_words

def load_products(num=None):
    """
    Loads products from the `items.json` file and combine them with reviews
    through `asin`.
    Return: dict[asin, product]
    """
    with open(ITEMS_PATH) as f:
        all_products = json.load(f)
        if num is not None:
            random.shuffle(all_products)
            all_products = all_products[:num]
        products = dict()
        asins = set()
        for p in all_products:
            asin = p['asin']
            if asin in asins:
                continue
            asins.add(asin)
            products[asin] = p

    with open(REVIEWS_PATH) as f:
        reviews = json.load(f)
        reviews = {r['asin']: r for r in reviews}

    for asin, p in products.items():
        if asin in reviews:
            p['review'] = reviews[asin]
        else:
            p['review'] = None
    return products

def get_top_attrs(attributes, k):
    attr_to_asins = defaultdict(list)

    for asin, attr_scores in attributes.items():
        top_attr_scoress = attr_scores[:k]
        for attr, score in top_attr_scoress:
            attr_to_asins[attr].append(asin)
    total = len([asin for asin, _ in attributes.items()])

    top_attrs = [
        (attr, len(asins) / total)
        for attr, asins in attr_to_asins.items()
    ]
    top_attrs = sorted(top_attrs, key=lambda x: -x[1])
    top_attrs = [f'{attr} | {score:.4f}' for attr, score in top_attrs]
    return top_attrs

def get_corpus(
        products,
        keys=('name', 'small_description'),
        category_type='category'
    ):
    """
    keys: `name`, `small_description`, `review`
    category_type: `category`, `query`
    """
    all_products = list(products.values())

    asins_by_cat = defaultdict(set)
    corpus_by_cat = defaultdict(list)
    for p in all_products:
        category = p[category_type]
        asin = p['asin']
        if asin in asins_by_cat[category]:
            continue
        asins_by_cat[category].add(asin)

        text = []
        for key in keys:
            if key == 'review':
                rs = p['review']['reviews']
                if r is not None:
                    text_ = ' '.join([r['review'].lower() for r in rs])
                else:
                    text_ = ''
            else:
                text_ = p[key].lower()
            text.append(text_)
        text = ' '.join(text)
        corpus_by_cat[category].append((asin, text))
    return corpus_by_cat

def generate_ngram_attrs(corpus_by_cat, ngram_range, k, attrs):
    vectorizer = TfidfVectorizer(
        stop_words=get_stop_words(),
        ngram_range=ngram_range,
        max_features=1000,
    )

    top_attrs_by_cat = dict()
    for category, corpus in tqdm(corpus_by_cat.items(),
                                 total=len(corpus_by_cat)):
        asins = [_[0] for _ in corpus]
        texts = [_[1] for _ in corpus]
        vec = vectorizer.fit_transform(texts).todense()
        df = pd.DataFrame(vec, columns=vectorizer.get_feature_names_out())

        attrs_by_cat = dict()
        for asin, (row_name, row) in zip(asins, df.iterrows()):
            attr_scores = sorted(
                list(zip(row.index, row)),
                key=lambda x: -x[1]
            )
            attrs_by_cat[asin] = attr_scores
            attrs[asin] = attr_scores
        top_attrs_by_cat[category.lower()] = get_top_attrs(attrs_by_cat, k=k)
    print(top_attrs_by_cat.keys())
    return top_attrs_by_cat

def generate_attrs(corpus_by_cat, k, save_name):
    attrs = dict()
    for n in range(1, 3):
        ngram_range = (n, n)
        top_attrs_by_cat = \
            generate_ngram_attrs(corpus_by_cat, ngram_range, k, attrs)

        if save_name is not None:
            save_path = Path(ATTR_DIR) / f'{save_name}_{n}-gram.yaml'
            with open(save_path, 'w') as f:
                yaml.dump(top_attrs_by_cat, f, default_flow_style=False)
            print(f'Saved: {save_path}')

    save_path = Path(ATTR_DIR) / f'{save_name}_attrs_unfiltered.json'
    with open(save_path, 'w') as f:
        json.dump(attrs, f)
    print(f'Saved: {save_path}')

if __name__ == '__main__':
    """
    python -m web_agent_site.attributes.generate_attrs

    Inspect in notebooks/attributes.ipynb.
    """
    products = load_products(num=40000)

    corpus_by_cat_broad = get_corpus(products, category_type='category')
    generate_attrs(corpus_by_cat_broad, k=5, save_name='broad')

    corpus_by_cat_narrow = get_corpus(products, category_type='query')
    generate_attrs(corpus_by_cat_narrow, k=5, save_name='narrow')


import json
import yaml
import random
from pathlib import Path
from collections import defaultdict

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_extraction import text as sk_text
import pandas as pd
from tqdm import tqdm
from rich import print

ITEMS_PATH = './data/ITEMS_mar1.json'
REVIEWS_PATH = './data/reviews.json'
ATTR_DIR = './data/attributes'

random.seed(0)

def get_stop_words():
    extra_stop_words = set([str(i) for i in range(1000)])
    stop_words = sk_text.ENGLISH_STOP_WORDS.union(extra_stop_words)
    return stop_words

def load_products(num=None):
    """
    Loads products from the `items.json` file and combine them with reviews
    through `asin`.
    Return: dict[asin, product]
    """
    with open(ITEMS_PATH) as f:
        all_products = json.load(f)
        if num is not None:
            random.shuffle(all_products)
            all_products = all_products[:num]
        products = dict()
        asins = set()
        for p in all_products:
            asin = p['asin']
            if asin in asins:
                continue
            asins.add(asin)


from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_extraction import text as sk_text
import pandas as pd
from tqdm import tqdm
from rich import print

ITEMS_PATH = './data/ITEMS_mar1.json'
REVIEWS_PATH = './data/reviews.json'
ATTR_DIR = './data/attributes'

random.seed(0)

def get_stop_words():
    extra_stop_words = set([str(i) for i in range(1000)])
    stop_words = sk_text.ENGLISH_STOP_WORDS.union(extra_stop_words)
    return stop_words

def load_products(num=None):
    """
    Loads products from the `items.json` file and combine them with reviews
    through `asin`.
    Return: dict[asin, product]
    """
    with open(ITEMS_PATH) as f:
        all_products = json.load(f)
        if num is not None:
            random.shuffle(all_products)
            all_products = all_products[:num]
        products = dict()
        asins = set()
        for p in all_products:
            asin = p['asin']
            if asin in asins:
                continue
            asins.add(asin)
            products[asin] = p

    with open(REVIEWS_PATH) as f:
        reviews = json.load(f)
        reviews = {r['asin']: r for r in reviews}

    for asin, p in products.items():
        if asin in reviews:
            p['review'] = reviews[asin]
        else:
            p['review'] = None
    return products

def get_top_attrs(attributes, k):
    attr_to_asins = defaultdict(list)

    for asin, attr_scores in attributes.items():
        top_attr_scoress = attr_scores[:k]
        for attr, score in top_attr_scoress:
            attr_to_asins[attr].append(asin)
    total = len([asin for asin, _ in attributes.items()])

    top_attrs = [
        (attr, len(asins) / total)
        for attr, asins in attr_to_asins.items()
    ]
    top_attrs = sorted(top_attrs, key=lambda x: -x[1])
    top_attrs = [f'{attr} | {score:.4f}' for attr, score in top_attrs]
    return top_attrs


        asins = set()
        for p in all_products:
            asin = p['asin']
            if asin in asins:
                continue
            asins.add(asin)
            products[asin] = p

    with open(REVIEWS_PATH) as f:
        reviews = json.load(f)
        reviews = {r['asin']: r for r in reviews}

    for asin, p in products.items():
        if asin in reviews:
            p['review'] = reviews[asin]
        else:
            p['review'] = None
    return products

def get_top_attrs(attributes, k):
    attr_to_asins = defaultdict(list)

    for asin, attr_scores in attributes.items():
        top_attr_scoress = attr_scores[:k]
        for attr, score in top_attr_scoress:
            attr_to_asins[attr].append(asin)
    total = len([asin for asin, _ in attributes.items()])

    top_attrs = [
        (attr, len(asins) / total)
        for attr, asins in attr_to_asins.items()
    ]
    top_attrs = sorted(top_attrs, key=lambda x: -x[1])
    top_attrs = [f'{attr} | {score:.4f}' for attr, score in top_attrs]
    return top_attrs

def get_corpus(
        products,
        keys=('name', 'small_description'),
        category_type='category'
    ):
    """
    keys: `name`, `small_description`, `review`
    category_type: `category`, `query`
    """
    all_products = list(products.values())

    asins_by_cat = defaultdict(set)
    corpus_by_cat = defaultdict(list)
    for p in all_products:
        category = p[category_type]
        asin = p['asin']
        if asin in asins_by_cat[category]:
            continue


def get_top_attrs(attributes, k):
    attr_to_asins = defaultdict(list)

    for asin, attr_scores in attributes.items():
        top_attr_scoress = attr_scores[:k]
        for attr, score in top_attr_scoress:
            attr_to_asins[attr].append(asin)
    total = len([asin for asin, _ in attributes.items()])

    top_attrs = [
        (attr, len(asins) / total)
        for attr, asins in attr_to_asins.items()
    ]
    top_attrs = sorted(top_attrs, key=lambda x: -x[1])
    top_attrs = [f'{attr} | {score:.4f}' for attr, score in top_attrs]
    return top_attrs

def get_corpus(
        products,
        keys=('name', 'small_description'),
        category_type='category'
    ):
    """
    keys: `name`, `small_description`, `review`
    category_type: `category`, `query`
    """
    all_products = list(products.values())

    asins_by_cat = defaultdict(set)
    corpus_by_cat = defaultdict(list)
    for p in all_products:
        category = p[category_type]
        asin = p['asin']
        if asin in asins_by_cat[category]:
            continue
        asins_by_cat[category].add(asin)

        text = []
        for key in keys:
            if key == 'review':
                rs = p['review']['reviews']
                if r is not None:
                    text_ = ' '.join([r['review'].lower() for r in rs])
                else:
                    text_ = ''
            else:
                text_ = p[key].lower()
            text.append(text_)
        text = ' '.join(text)
        corpus_by_cat[category].append((asin, text))
    return corpus_by_cat

def generate_ngram_attrs(corpus_by_cat, ngram_range, k, attrs):
    vectorizer = TfidfVectorizer(
        stop_words=get_stop_words(),
        ngram_range=ngram_range,
        max_features=1000,
    )

    top_attrs_by_cat = dict()
    for category, corpus in tqdm(corpus_by_cat.items(),
                                 total=len(corpus_by_cat)):
        asins = [_[0] for _ in corpus]
        texts = [_[1] for _ in corpus]
        vec = vectorizer.fit_transform(texts).todense()
        df = pd.DataFrame(vec, columns=vectorizer.get_feature_names_out())

        attrs_by_cat = dict()
        for asin, (row_name, row) in zip(asins, df.iterrows()):
            attr_scores = sorted(


        if asin in asins_by_cat[category]:
            continue
        asins_by_cat[category].add(asin)

        text = []
        for key in keys:
            if key == 'review':
                rs = p['review']['reviews']
                if r is not None:
                    text_ = ' '.join([r['review'].lower() for r in rs])
                else:
                    text_ = ''
            else:
                text_ = p[key].lower()
            text.append(text_)
        text = ' '.join(text)
        corpus_by_cat[category].append((asin, text))
    return corpus_by_cat

def generate_ngram_attrs(corpus_by_cat, ngram_range, k, attrs):
    vectorizer = TfidfVectorizer(
        stop_words=get_stop_words(),
        ngram_range=ngram_range,
        max_features=1000,
    )

    top_attrs_by_cat = dict()
    for category, corpus in tqdm(corpus_by_cat.items(),
                                 total=len(corpus_by_cat)):
        asins = [_[0] for _ in corpus]
        texts = [_[1] for _ in corpus]
        vec = vectorizer.fit_transform(texts).todense()
        df = pd.DataFrame(vec, columns=vectorizer.get_feature_names_out())

        attrs_by_cat = dict()
        for asin, (row_name, row) in zip(asins, df.iterrows()):
            attr_scores = sorted(
                list(zip(row.index, row)),
                key=lambda x: -x[1]
            )
            attrs_by_cat[asin] = attr_scores
            attrs[asin] = attr_scores
        top_attrs_by_cat[category.lower()] = get_top_attrs(attrs_by_cat, k=k)
    print(top_attrs_by_cat.keys())
    return top_attrs_by_cat

def generate_attrs(corpus_by_cat, k, save_name):
    attrs = dict()
    for n in range(1, 3):
        ngram_range = (n, n)
        top_attrs_by_cat = \
            generate_ngram_attrs(corpus_by_cat, ngram_range, k, attrs)

        if save_name is not None:
            save_path = Path(ATTR_DIR) / f'{save_name}_{n}-gram.yaml'
            with open(save_path, 'w') as f:
                yaml.dump(top_attrs_by_cat, f, default_flow_style=False)
            print(f'Saved: {save_path}')

    save_path = Path(ATTR_DIR) / f'{save_name}_attrs_unfiltered.json'
    with open(save_path, 'w') as f:
        json.dump(attrs, f)
    print(f'Saved: {save_path}')


    for category, corpus in tqdm(corpus_by_cat.items(),
                                 total=len(corpus_by_cat)):
        asins = [_[0] for _ in corpus]
        texts = [_[1] for _ in corpus]
        vec = vectorizer.fit_transform(texts).todense()
        df = pd.DataFrame(vec, columns=vectorizer.get_feature_names_out())

        attrs_by_cat = dict()
        for asin, (row_name, row) in zip(asins, df.iterrows()):
            attr_scores = sorted(
                list(zip(row.index, row)),
                key=lambda x: -x[1]
            )
            attrs_by_cat[asin] = attr_scores
            attrs[asin] = attr_scores
        top_attrs_by_cat[category.lower()] = get_top_attrs(attrs_by_cat, k=k)
    print(top_attrs_by_cat.keys())
    return top_attrs_by_cat

def generate_attrs(corpus_by_cat, k, save_name):
    attrs = dict()
    for n in range(1, 3):
        ngram_range = (n, n)
        top_attrs_by_cat = \
            generate_ngram_attrs(corpus_by_cat, ngram_range, k, attrs)

        if save_name is not None:
            save_path = Path(ATTR_DIR) / f'{save_name}_{n}-gram.yaml'
            with open(save_path, 'w') as f:
                yaml.dump(top_attrs_by_cat, f, default_flow_style=False)
            print(f'Saved: {save_path}')

    save_path = Path(ATTR_DIR) / f'{save_name}_attrs_unfiltered.json'
    with open(save_path, 'w') as f:
        json.dump(attrs, f)
    print(f'Saved: {save_path}')

if __name__ == '__main__':
    """
    python -m web_agent_site.attributes.generate_attrs

    Inspect in notebooks/attributes.ipynb.
    """
    products = load_products(num=40000)

    corpus_by_cat_broad = get_corpus(products, category_type='category')
    generate_attrs(corpus_by_cat_broad, k=5, save_name='broad')

    corpus_by_cat_narrow = get_corpus(products, category_type='query')
    generate_attrs(corpus_by_cat_narrow, k=5, save_name='narrow')


import yaml
from pathlib import Path
from rich import print

ATTR_DIR = './data/attributes'

ATTR_PATHS = [
    'narrow_2-gram.yaml',
    'narrow_1-gram.yaml',
    'broad_2-gram.yaml',
    'broad_1-gram.yaml',
]
ATTR_PATHS = [Path(ATTR_DIR) / af for af in ATTR_PATHS]

def annotate(attr_path):
    with open(attr_path) as f:
        attrs_by_cat = yaml.safe_load(f)

    unique_attrs = set()
    all_attrs = []
    for _, attrs in attrs_by_cat.items():
        attrs = [a.split('|')[0].strip() for a in attrs]
        unique_attrs.update(attrs)
        all_attrs += attrs
    print(f'Total unique attributes: {len(unique_attrs)}')
    total = len(all_attrs)
    num_left = len(all_attrs)

    annotated_attrs_by_cat = dict()
    for category, attrs in attrs_by_cat.items():
        print(
            f'Category: [ {category} ] | '
            f'Number of attributes: {len(attrs)}\n'
        )
        annotated_attrs = []
        for i, attr in enumerate(attrs):
            attr, score = attr.split(' | ')
            print(
                f'{"[" + str(i) + "]":<5} '
                f'[bold green]{attr:<30}[/bold green] | '
                f'[red]{category}[/red] | '
                f'{score}'
            )
            tags = input(
                'Annotate [1: ITEM, 2: PROP, 3: USE, '
                '⎵: next example, q: next category] > '
            )
            print('\n')
            tags = tags.strip()
            annotated_attrs.append(f'{attr} | {score} | {tags}')
            if 'q' in tags:
                break

        num_left -= len(attrs)
        print(f'{num_left} / {total} total attributes left.')

        ans = input('Starting the next category... [y/n] > ')
        if ans == 'n':
            break

def main():
    for attr_path in ATTR_PATHS:
        annotate(attr_path)

if __name__ == '__main__':
    """
    python -m web_agent_site.attributes.annotate
    """
    main()


main

                f'[red]{category}[/red] | '
                f'{score}'
            )
            tags = input(
                'Annotate [1: ITEM, 2: PROP, 3: USE, '
                '⎵: next example, q: next category] > '
            )
            print('\n')
            tags = tags.strip()
            annotated_attrs.append(f'{attr} | {score} | {tags}')
            if 'q' in tags:
                break

        num_left -= len(attrs)
        print(f'{num_left} / {total} total attributes left.')

        ans = input('Starting the next category... [y/n] > ')
        if ans == 'n':
            break

def main():
    for attr_path in ATTR_PATHS:
        annotate(attr_path)

if __name__ == '__main__':
    """
    python -m web_agent_site.attributes.annotate
    """
    main()


import gym
import random
import requests
import string
import time

from bs4 import BeautifulSoup
from bs4.element import Comment
from gym import spaces
from os.path import join, dirname, abspath
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.keys import Keys
from selenium.common.exceptions import ElementNotInteractableException
from web_agent_site.engine.engine import parse_action, END_BUTTON

class WebAgentSiteEnv(gym.Env):
    """Gym environment for HTML mode of WebShop environment"""

    def __init__(self, observation_mode='html', **kwargs):
        """
        Constructor for HTML environment

        Arguments:
        observation_mode (`str`) -- ['html' | 'text'] (default 'html')
        pause (`float`) -- Pause (in seconds) after taking an action.
            This is mainly for demo purposes.
            Recommended value: 2.0s
        render (`bool`) -- Show browser if set to `True`.
        session ('str') -- Session ID to initialize environment with
        """
        super(WebAgentSiteEnv, self).__init__()
        self.observation_mode = observation_mode
        self.kwargs = kwargs

        # Create a browser driver to simulate the WebShop site
        service = Service(join(dirname(abspath(__file__)), 'chromedriver'))
        options = Options()
        if 'render' not in kwargs or not kwargs['render']:
            options.add_argument("--headless")  # don't show browser
        self.browser = webdriver.Chrome(service=service, options=options)

        # Set flags and values for WebShop session
        self.text_to_clickable = None
        self.assigned_session = kwargs.get('session')
        self.session = None
        self.reset()

    def step(self, action):
        """
        Takes an action, updates WebShop environment, and returns (observation, reward, done, info)

        Arguments:
        action (`str`): An action should be of the following structure:
          - search[keywords]
          - click[value]
        If action not valid, perform nothing.
        """
        reward = 0.0
        done = False
        info = None

        # Map action to executed command on the WebShop environment via the broswer driver
        action_name, action_arg = parse_action(action)
        if action_name == 'search':
            try:
                search_bar = self.browser.find_element_by_id('search_input')
            except Exception:
                pass
            else:
                search_bar.send_keys(action_arg)
                search_bar.submit()
        elif action_name == 'click':
            try:
                self.text_to_clickable[action_arg].click()
            except ElementNotInteractableException:
                # Perform force click with JavaScript
                button = self.text_to_clickable[action_arg]
                self.browser.execute_script("arguments[0].click();", button)
            reward = self.get_reward()
            if action_arg == END_BUTTON:
                done = True
        elif action_name == 'end':
            done = True
        else:
            print('Invalid action. No action performed.')

        if 'pause' in self.kwargs:
            time.sleep(self.kwargs['pause'])
        return self.observation, reward, done, info

    def get_available_actions(self):
        """Returns list of available actions at the current step"""
        # Determine if a search bar is available
        try:
            search_bar = self.browser.find_element_by_id('search_input')
        except Exception:
            has_search_bar = False
        else:
            has_search_bar = True

        # Collect buttons, links, and options as clickables
        buttons = self.browser.find_elements_by_class_name('btn')
        product_links = self.browser.find_elements_by_class_name('product-link')
        buying_options = self.browser.find_elements_by_css_selector("input[type='radio']")

        self.text_to_clickable = {
            f'{b.text}': b
            for b in buttons + product_links
        }
        for opt in buying_options:
            opt_value = opt.get_attribute('value')
            self.text_to_clickable[f'{opt_value}'] = opt
        return dict(
            has_search_bar=has_search_bar,
            clickables=list(self.text_to_clickable.keys()),
        )

    def _parse_html(self, html=None, url=None):
        """
        Returns web request result wrapped in BeautifulSoup object

        Arguments:
        url (`str`): If no url or html is provided, use the current
            observation (HTML) for parsing.
        """
        if html is None:
            if url is not None:
                html = requests.get(url)
            else:
                html = self.state['html']
        html_obj = BeautifulSoup(html, 'html.parser')
        return html_obj

    def get_reward(self):
        """Get reward value at current step of the environment"""
        html_obj = self._parse_html()
        r = html_obj.find(id='reward')
        r = float(r.findChildren("pre")[0].string) if r is not None else 0.0
        return r

    def get_instruction_text(self):
        """Get corresponding instruction text for environment current step"""
        html_obj = self._parse_html(self.browser.page_source)
        instruction_text = html_obj.find(id='instruction-text').h4.text
        return instruction_text

    def convert_html_to_text(self, html):
        """Strip HTML of tags and add separators to convert observation into simple mode"""
        texts = self._parse_html(html).findAll(text=True)
        visible_texts = filter(tag_visible, texts)
        observation = ' [SEP] '.join(t.strip() for t in visible_texts if t != '\n')
        return observation

    @property
    def state(self):
        """
        State that includes all information. The actual observation are
        likely to be a subset or reduced form of the state.
        """
        return dict(
            url=self.browser.current_url,
            html=self.browser.page_source,
            instruction_text=self.instruction_text,
        )

    @property
    def observation(self):
        """Compiles state into either the `html` or `text` observation mode"""
        html = self.state['html']
        if self.observation_mode == 'html':
            return html
        elif self.observation_mode == 'text':
            return self.convert_html_to_text(html)
        else:
            raise ValueError(
                f'Observation mode {self.observation_mode} not supported.'
            )

    @property
    def action_space(self):
        # Recommended to use `get_available_actions` instead
        return NotImplementedError

    @property
    def observation_space(self):
        return NotImplementedError

    def reset(self):
        """Create a new session and reset environment variables"""
        if self.assigned_session is not None:
            self.session = self.assigned_session
        else:
            self.session = ''.join(random.choices(string.ascii_lowercase, k=5))
        init_url = f'http://127.0.0.1:3000/{self.session}'
        self.browser.get(init_url)

        self.instruction_text = self.get_instruction_text()

        return self.observation, None

    def render(self, mode='human'):
        # TODO: Render observation in terminal or WebShop website
        return NotImplementedError

    def close(self):
        # TODO: When DB used instead of JSONs, tear down DB here
        self.browser.close()
        print('Browser closed.')

def tag_visible(element):
    """Helper method to strip HTML block of extraneous tags"""
    ignore = {'style', 'script', 'head', 'title', 'meta', '[document]'}
    return (
        element.parent.name not in ignore and not isinstance(element, Comment)
    )


        if self.assigned_session is not None:
            self.session = self.assigned_session
        else:
            self.session = ''.join(random.choices(string.ascii_lowercase, k=5))
        init_url = f'http://127.0.0.1:3000/{self.session}'
        self.browser.get(init_url)

        self.instruction_text = self.get_instruction_text()

        return self.observation, None

    def render(self, mode='human'):
        # TODO: Render observation in terminal or WebShop website
        return NotImplementedError

    def close(self):
        # TODO: When DB used instead of JSONs, tear down DB here
        self.browser.close()
        print('Browser closed.')

def tag_visible(element):
    """Helper method to strip HTML block of extraneous tags"""
    ignore = {'style', 'script', 'head', 'title', 'meta', '[document]'}
    return (
        element.parent.name not in ignore and not isinstance(element, Comment)
    )